Adding option to re-raise error when loading a spectrum

.codecov.yml

@@ -12,7 +12,7 @@ coverage:
     patch:
       # just the lines changed
       default:
-        target: 50%
+        target: 30%
         threshold: 10%
         base: auto
         only_pulls: false

astrodbkit2/spectra.py

@@ -181,8 +181,21 @@ def wcs1d_multispec_loader(file_obj, flux_unit=None,
                       meta=meta)
 
 
-def load_spectrum(filename, spectra_format=None):
-    """Attempt to load the filename as a spectrum object"""
+def load_spectrum(filename: str,
+                  spectra_format: str=None,
+                  raise_error: bool=False):
+    """Attempt to load the filename as a spectrum object
+
+    Parameters
+    ----------
+    filename
+        Name of the file to read
+    spectra_format
+        Optional file format, passed to Spectrum1D.read. 
+        In its absense Spectrum1D.read will attempt to determine the format.
+    raise_error
+        Boolean to control if a failure to read the spectrum should raise an error.
+    """
 
     # Convert filename if using environment variables
     if filename.startswith('$'):
@@ -201,7 +214,13 @@ def load_spectrum(filename, spectra_format=None):
         else:
             spec1d = Spectrum1D.read(filename)
     except Exception as e:  # pylint: disable=broad-except, invalid-name
-        print(f'Error loading {filename}: {e}')
-        spec1d = filename
+        msg = f'Error loading {filename}: {e}'
+
+        # Control whether an error should be explicitly raised if failing to read
+        if raise_error:
+            raise TypeError(msg) from e
+        else:
+            print(msg)
+            spec1d = filename
 
     return spec1d

astrodbkit2/tests/test_spectra.py

@@ -4,7 +4,15 @@
 import numpy as np
 from astropy.io import fits
 from astropy.units import Unit
-from astrodbkit2.spectra import identify_spex_prism, _identify_spex, load_spectrum, spex_prism_loader, identify_wcs1d_multispec, wcs1d_multispec_loader
+from astrodbkit2.spectra import (
+    identify_spex_prism,
+    _identify_spex,
+    load_spectrum,
+    spex_prism_loader,
+    identify_wcs1d_multispec,
+    wcs1d_multispec_loader,
+)
+
 try:
     import mock
 except ImportError:
@@ -16,11 +24,11 @@ def good_spex_file():
     fake_data = np.arange(1, 10, 1)
     n = np.array([fake_data, fake_data, fake_data])
     hdr = fits.Header()
-    hdr['TELESCOP'] = 'NASA IRTF'
-    hdr['INSTRUME'] = 'SPeX, IRTF Spectrograph'
-    hdr['GRAT'] = 'LowRes15 '
-    hdr['XUNITS'] = 'Microns '
-    hdr['YUNITS'] = 'ergs s-1 cm-2 A-1'
+    hdr["TELESCOP"] = "NASA IRTF"
+    hdr["INSTRUME"] = "SPeX, IRTF Spectrograph"
+    hdr["GRAT"] = "LowRes15 "
+    hdr["XUNITS"] = "Microns "
+    hdr["YUNITS"] = "ergs s-1 cm-2 A-1"
     hdu1 = fits.PrimaryHDU(n, header=hdr)
     return fits.HDUList([hdu1])
 
@@ -30,10 +38,10 @@ def bad_spex_file():
     fake_data = np.arange(1, 10, 1)
     n = np.array([fake_data, fake_data, fake_data])
     hdr = fits.Header()
-    hdr['TELESCOP'] = 'MISSING'
-    hdr['INSTRUME'] = 'MISSING'
-    hdr['GRAT'] = 'MISSING'
-    hdr['XUNITS'] = 'UNKNOWN'
+    hdr["TELESCOP"] = "MISSING"
+    hdr["INSTRUME"] = "MISSING"
+    hdr["GRAT"] = "MISSING"
+    hdr["XUNITS"] = "UNKNOWN"
     hdu1 = fits.PrimaryHDU(n, header=hdr)
     return fits.HDUList([hdu1])
 
@@ -42,104 +50,111 @@ def bad_spex_file():
 def good_wcs1dmultispec():
     n = np.empty((2141, 1, 4))
     hdr = fits.Header()
-    hdr['WCSDIM'] = 3
-    hdr['NAXIS'] = 3
-    hdr['WAT0_001'] = 'system=equispec'
-    hdr['WAT1_001'] = 'wtype=linear label=Wavelength units=angstroms'
-    hdr['WAT2_001'] = 'wtype=linear'
-    hdr['BANDID1'] = 'spectrum - background fit, weights variance, clean no'               
-    hdr['BANDID2'] = 'raw - background fit, weights none, clean no'                        
-    hdr['BANDID3'] = 'background - background fit'                                         
-    hdr['BANDID4'] = 'sigma - background fit, weights variance, clean no'  
-    hdr['CTYPE1'] = 'LINEAR  '
-    hdr['BUNIT']   = 'erg/cm2/s/A'
+    hdr["WCSDIM"] = 3
+    hdr["NAXIS"] = 3
+    hdr["WAT0_001"] = "system=equispec"
+    hdr["WAT1_001"] = "wtype=linear label=Wavelength units=angstroms"
+    hdr["WAT2_001"] = "wtype=linear"
+    hdr["BANDID1"] = "spectrum - background fit, weights variance, clean no"
+    hdr["BANDID2"] = "raw - background fit, weights none, clean no"
+    hdr["BANDID3"] = "background - background fit"
+    hdr["BANDID4"] = "sigma - background fit, weights variance, clean no"
+    hdr["CTYPE1"] = "LINEAR  "
+    hdr["BUNIT"] = "erg/cm2/s/A"
     hdu1 = fits.PrimaryHDU(n, header=hdr)
     return fits.HDUList([hdu1])
 
+
 @pytest.fixture(scope="module")
 def alt_wcs1dmultispec():
     n = np.empty((2141,))
     hdr = fits.Header()
-    hdr['WCSDIM'] = 1
-    hdr['NAXIS'] = 1
-    hdr['WAT0_001'] = 'system=equispec'
-    hdr['WAT1_001'] = 'wtype=linear label=Wavelength units=angstroms'
-    hdr['WAT2_001'] = 'wtype=linear'
-    hdr['BANDID1'] = 'spectrum - background fit, weights variance, clean no'               
-    hdr['BANDID2'] = 'raw - background fit, weights none, clean no'                        
-    hdr['BANDID3'] = 'background - background fit'                                         
-    hdr['BANDID4'] = 'sigma - background fit, weights variance, clean no'  
-    hdr['CTYPE1'] = 'LINEAR  '
-    hdr['BUNIT']   = 'erg/cm2/s/A'
+    hdr["WCSDIM"] = 1
+    hdr["NAXIS"] = 1
+    hdr["WAT0_001"] = "system=equispec"
+    hdr["WAT1_001"] = "wtype=linear label=Wavelength units=angstroms"
+    hdr["WAT2_001"] = "wtype=linear"
+    hdr["BANDID1"] = "spectrum - background fit, weights variance, clean no"
+    hdr["BANDID2"] = "raw - background fit, weights none, clean no"
+    hdr["BANDID3"] = "background - background fit"
+    hdr["BANDID4"] = "sigma - background fit, weights variance, clean no"
+    hdr["CTYPE1"] = "LINEAR  "
+    hdr["BUNIT"] = "erg/cm2/s/A"
     hdu1 = fits.PrimaryHDU(n, header=hdr)
     return fits.HDUList([hdu1])
 
 
-@mock.patch('astrodbkit2.spectra.fits.open')
+@mock.patch("astrodbkit2.spectra.fits.open")
 def test_identify_spex_prism(mock_fits_open, good_spex_file):
     mock_fits_open.return_value = good_spex_file
 
-    filename = 'https://s3.amazonaws.com/bdnyc/SpeX/Prism/U10013_SpeX.fits'
-    assert identify_spex_prism('read', filename)
-    filename = 'I am not a valid spex prism file'
-    assert not identify_spex_prism('read', filename)
+    filename = "https://s3.amazonaws.com/bdnyc/SpeX/Prism/U10013_SpeX.fits"
+    assert identify_spex_prism("read", filename)
+    filename = "I am not a valid spex prism file"
+    assert not identify_spex_prism("read", filename)
 
 
-@mock.patch('astrodbkit2.spectra.fits.open')
+@mock.patch("astrodbkit2.spectra.fits.open")
 def test_identify_spex(mock_fits_open, good_spex_file, bad_spex_file):
     mock_fits_open.return_value = good_spex_file
-    assert _identify_spex('filename')
+    assert _identify_spex("filename")
     mock_fits_open.return_value = bad_spex_file
-    assert not _identify_spex('filename')
+    assert not _identify_spex("filename")
 
 
-@mock.patch('astrodbkit2.spectra.fits.open')
+@mock.patch("astrodbkit2.spectra.fits.open")
 def test_load_spex_prism(mock_fits_open, good_spex_file, bad_spex_file):
     # Test good example
     mock_fits_open.return_value = good_spex_file
-    spectrum = spex_prism_loader('filename')
-    assert spectrum.unit == Unit('erg / (A cm2 s)')
+    spectrum = spex_prism_loader("filename")
+    assert spectrum.unit == Unit("erg / (A cm2 s)")
     # Test bad example
     mock_fits_open.return_value = bad_spex_file
-    spectrum = spex_prism_loader('filename')
-    assert spectrum.unit == Unit('erg')
+    spectrum = spex_prism_loader("filename")
+    assert spectrum.unit == Unit("erg")
 
 
-@mock.patch('astrodbkit2.spectra.read_fileobj_or_hdulist')
+@mock.patch("astrodbkit2.spectra.read_fileobj_or_hdulist")
 def test_identify_wcs1d_multispec(mock_fits_open, good_wcs1dmultispec):
     mock_fits_open.return_value = good_wcs1dmultispec
 
-    filename = 'https://s3.amazonaws.com/bdnyc/optical_spectra/U10929.fits'
-    assert identify_wcs1d_multispec('read', filename)
+    filename = "https://s3.amazonaws.com/bdnyc/optical_spectra/U10929.fits"
+    assert identify_wcs1d_multispec("read", filename)
 
 
-@mock.patch('astrodbkit2.spectra.read_fileobj_or_hdulist')
+@mock.patch("astrodbkit2.spectra.read_fileobj_or_hdulist")
 def test_wcs1d_multispec_loader(mock_fits_open, good_wcs1dmultispec, alt_wcs1dmultispec):
     mock_fits_open.return_value = good_wcs1dmultispec
 
-    spectrum = wcs1d_multispec_loader('filename')
-    assert spectrum.unit == Unit('erg / (Angstrom cm2 s)')
-    assert spectrum.wavelength.unit == Unit('Angstrom')
+    spectrum = wcs1d_multispec_loader("filename")
+    assert spectrum.unit == Unit("erg / (Angstrom cm2 s)")
+    assert spectrum.wavelength.unit == Unit("Angstrom")
 
     # Check flux_unit is converted correctly
-    spectrum = wcs1d_multispec_loader('filename', flux_unit=Unit('erg / (um cm2 s)'))
-    assert spectrum.unit == Unit('erg / (um cm2 s)')
+    spectrum = wcs1d_multispec_loader("filename", flux_unit=Unit("erg / (um cm2 s)"))
+    assert spectrum.unit == Unit("erg / (um cm2 s)")
 
     # NAXIS=1 example
     mock_fits_open.return_value = alt_wcs1dmultispec
-    spectrum = wcs1d_multispec_loader('filename')
-    assert spectrum.unit == Unit('erg / (Angstrom cm2 s)')
-    assert spectrum.wavelength.unit == Unit('Angstrom')
+    spectrum = wcs1d_multispec_loader("filename")
+    assert spectrum.unit == Unit("erg / (Angstrom cm2 s)")
+    assert spectrum.wavelength.unit == Unit("Angstrom")
 
 
-@mock.patch('astrodbkit2.spectra.Spectrum1D.read')
+@mock.patch("astrodbkit2.spectra.Spectrum1D.read")
 def test_load_spectrum(mock_spectrum1d, monkeypatch):
-    _ = load_spectrum('fake_file.txt')
-    mock_spectrum1d.assert_called_with('fake_file.txt')
-    _ = load_spectrum('fake_file.txt', spectra_format='SpeX')
-    mock_spectrum1d.assert_called_with('fake_file.txt', format='SpeX')
+    _ = load_spectrum("fake_file.txt")
+    mock_spectrum1d.assert_called_with("fake_file.txt")
+    _ = load_spectrum("fake_file.txt", spectra_format="SpeX")
+    mock_spectrum1d.assert_called_with("fake_file.txt", format="SpeX")
 
     # Testing user-set environment variable
-    monkeypatch.setenv('FAKE_ENV', '/User/path')
-    _ = load_spectrum('$FAKE_ENV/to/my/fake_file.txt')
-    mock_spectrum1d.assert_called_with('/User/path/to/my/fake_file.txt')
+    monkeypatch.setenv("FAKE_ENV", "/User/path")
+    _ = load_spectrum("$FAKE_ENV/to/my/fake_file.txt")
+    mock_spectrum1d.assert_called_with("/User/path/to/my/fake_file.txt")
+
+
+def test_load_spectrum_error():
+    # Test error handling
+    with pytest.raises(TypeError):
+        _ = load_spectrum("fake_file.fits", raise_error=True)