JP-3753: Refactor extract1d

jwst/datamodels/container.py

@@ -330,7 +330,7 @@ def save(self,
         path : str or None
             - If None, the `meta.filename` is used for each model.
             - If a string, the string is used as a root and an index is
-              appended.
+              appended, along with the '.fits' extension.
 
         save_model_func: func or None
             Alternate function to save each model instead of
@@ -353,8 +353,11 @@ def save(self,
                     save_path = model.meta.filename
                 else:
                     if len(self) <= 1:
-                        idx = None
-                    save_path = path+str(idx)+".fits"
+                        idx = ''
+                    if path.endswith(".fits"):
+                        save_path = path.replace(".fits", f"{idx}.fits")
+                    else:
+                        save_path = f"{path}{idx}.fits"
                 output_paths.append(
                     model.save(save_path, **kwargs)
                     )

jwst/datamodels/tests/test_model_container.py

@@ -127,7 +127,8 @@ def test_group_id(tmp_path):
     assert asn_group_ids == model_droup_ids
 
 
-def test_save(tmp_cwd, container):
+@pytest.mark.parametrize("path", ["foo", "foo.fits"])
+def test_save(tmp_cwd, container, path):
 
     # container pushes us to data/ directory so need to go back to tmp_cwd
     # to avoid polluting the data/ directory
@@ -142,8 +143,13 @@ def test_save(tmp_cwd, container):
             assert os.path.exists(fname)
 
         # test specifying path saves to custom path with indices
-        path = "foo"
         container.save(path)
-        expected_fnames = [path+str(i)+".fits" for i in range(len(container))]
+        expected_fnames = [path.replace(".fits", "")+str(i)+".fits" for i in range(len(container))]
         for fname in expected_fnames:
             assert os.path.exists(fname)
+
+        # test saving path when the container has length 1: no index appended
+        container = ModelContainer([container[0]])
+        container.save(path)
+        expected_fname = path.replace(".fits", "") + ".fits"
+        assert os.path.exists(expected_fname)

jwst/extract_1d/apply_apcorr.py

@@ -53,7 +53,7 @@ class ApCorrBase(abc.ABC):
     }
 
     def __init__(self, input_model, apcorr_table, sizeunit,
-                 location = None, slit_name = None, **match_kwargs):
+                 location=None, slit_name=None, **match_kwargs):
         self.correction = None
 
         self.model = input_model
@@ -155,7 +155,7 @@ def apply(self, spec_table):
         """
         flux_cols_to_correct = ('flux', 'flux_error', 'surf_bright', 'sb_error')
         var_cols_to_correct = ('flux_var_poisson', 'flux_var_rnoise', 'flux_var_flat', 
-                            'sb_var_poisson', 'sb_var_rnoise', 'sb_var_flat')
+                               'sb_var_poisson', 'sb_var_rnoise', 'sb_var_flat')
 
         for row in spec_table:
             correction = self.apcorr_func(row['npixels'], row['wavelength'])
@@ -180,7 +180,7 @@ class ApCorrPhase(ApCorrBase):
     """
     size_key = 'size'
 
-    def __init__(self, *args, pixphase = 0.5, **kwargs):
+    def __init__(self, *args, pixphase=0.5, **kwargs):
         self.phase = pixphase  # In the future we'll attempt to measure the pixel phase from inputs.
 
         super().__init__(*args, **kwargs)
@@ -297,7 +297,7 @@ class ApCorrRadial(ApCorrBase):
     """Aperture correction class used with spectral data produced from an extraction aperture radius."""
 
     def __init__(self, input_model, apcorr_table,
-                 location = None):
+                 location=None):
 
         self.correction = None
         self.model = input_model
@@ -339,7 +339,7 @@ def apply(self, spec_table):
         """
         flux_cols_to_correct = ('flux', 'flux_error', 'surf_bright', 'sb_error')
         var_cols_to_correct = ('flux_var_poisson', 'flux_var_rnoise', 'flux_var_flat', 
-                            'sb_var_poisson', 'sb_var_rnoise', 'sb_var_flat')
+                               'sb_var_poisson', 'sb_var_rnoise', 'sb_var_flat')
 
         for i, row in enumerate(spec_table):
             correction = self.apcorr_correction[i]
@@ -377,7 +377,7 @@ def match_wavelengths(self, wavelength_ifu):
     def find_apcorr_func(self, iwave, radius_ifu):
         # at ifu wavelength plane (iwave), the extraction radius is radius_ifu
         # pull out the radius values (self.size)  to use in the apcor ref file for this iwave
-        # self.size and self.apcorr have already been interpolated in wavelength to match the
+        # self.size and self.apcorr have already been interpolated in wavelength to match
         # the ifu wavelength range.
 
         radius_apcor = self.size[:, iwave]