simplifying query_step_status function

jwst/master_background/master_background_mos_step.py

@@ -107,7 +107,8 @@ def process(self, data):
             # If some type of background processing had already been done. Abort.
             # UNLESS forcing is enacted.
             if not self.force_subtract and \
-               'COMPLETE' in [query_step_status(data_model, "back_sub"), query_step_status(data_model, "master_background")]:
+               'COMPLETE' in [query_step_status_container(data_model, "back_sub"), \
+                              query_step_status_container(data_model, "master_background")]:
                 self.log.info('Background subtraction has already occurred. Skipping.')
                 record_step_status(data, 'master_background', success=False)
                 return data
@@ -275,3 +276,27 @@ def _calc_master_background(self, data, user_background=None):
             mb_multislit = self.flat_field(mb_multislit)
 
         return master_background, mb_multislit
+
+
+def query_step_status_container(model: datamodels.JwstDataModel, step_name: str):
+    '''Query the status of a step in a datamodel or container of datamodels
+    assuming the step status for the zeroth model is representative of all models.
+    This is true for the master background step but may not be generally applicable.
+    
+    Parameters
+    ----------
+    model : `~jwst.datamodels.JwstDataModel` instance
+        This is the datamodel or container of datamodels to check
+        
+    step_name : str
+        The attribute in meta.cal_step to check
+        
+    Returns
+    -------
+    status : str
+        The status of the step in meta.cal_step, typically 'COMPLETE' or 'SKIPPED'
+    '''
+    if isinstance(model, datamodels.ModelContainer):
+        return query_step_status(model[0], step_name)
+    else:
+        return query_step_status(model, step_name)

jwst/pipeline/calwebb_spec2.py

@@ -333,7 +333,7 @@ def process_exposure_product(
             # as DO_NOT_USE or NON_SCIENCE.
             resampled = self.pixel_replace(resampled)
             resampled = self.cube_build(resampled)
-            if query_step_status(resampled, "cube_build") == 'COMPLETE':
+            if query_step_status(resampled[0], "cube_build") == 'COMPLETE':
                 self.save_model(resampled[0], suffix='s3d')
         elif exp_type in ['MIR_LRS-SLITLESS']:
             resampled = calibrated.copy()
@@ -347,7 +347,7 @@ def process_exposure_product(
             # as DO_NOT_USE or NON_SCIENCE.
             resampled = self.pixel_replace(resampled)
         # Extract a 1D spectrum from the 2D/3D data
-        if exp_type in ['MIR_MRS', 'NRS_IFU'] and query_step_status(resampled, "cube_build") == 'SKIPPED':
+        if exp_type in ['MIR_MRS', 'NRS_IFU'] and query_step_status(resampled[0], "cube_build") == 'SKIPPED':
             # Skip extract_1d for IFU modes where no cube was built
             self.extract_1d.skip = True
 

jwst/stpipe/utilities.py

@@ -36,7 +36,6 @@
 import os
 import re
 from collections.abc import Sequence
-import numpy as np
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -166,9 +165,9 @@ def record_step_status(datamodel, cal_step, success=True):
         If True, then 'COMPLETE' is recorded.  If False, then 'SKIPPED'
     """
     if success:
-        status = 'COMPLETE'
+        status = COMPLETE
     else:
-        status = 'SKIPPED'
+        status = SKIPPED
 
     if isinstance(datamodel, Sequence):
         for model in datamodel:
@@ -195,17 +194,4 @@ def query_step_status(datamodel, cal_step):
     status : str
         The status of the step in meta.cal_step, typically 'COMPLETE' or 'SKIPPED'
     """
-    if isinstance(datamodel, Sequence):
-        status = np.array([getattr(model.meta.cal_step, cal_step, NOT_SET) for model in datamodel])
-        if np.any(status == COMPLETE):
-            return COMPLETE
-        elif np.all(status == NOT_SET):
-            return NOT_SET
-        elif np.all(status == SKIPPED):
-            return SKIPPED
-        else:
-            msg = f"Unrecognized status in meta.cal_step.{cal_step}."
-            raise ValueError(msg)
-        return getattr(datamodel[0].meta.cal_step, cal_step, NOT_SET)
-    else:
-        return getattr(datamodel.meta.cal_step, cal_step, NOT_SET)
+    return getattr(datamodel.meta.cal_step, cal_step, NOT_SET)