JP-3721: Simplify ModelContainer

changes/190.misc.rst

@@ -0,0 +1 @@
+improve support for list-like input into Step

src/stpipe/step.py

@@ -491,31 +491,37 @@
                                         e,
                                     )
                                 library.shelve(model, i)
-                    elif isinstance(args[0], AbstractDataModel):
-                        if self.class_alias is not None:
-                            if isinstance(args[0], Sequence):
-                                for model in args[0]:
-                                    try:
-                                        model[f"meta.cal_step.{self.class_alias}"] = (
-                                            "SKIPPED"
-                                        )
-                                    except AttributeError as e:  # noqa: PERF203
-                                        self.log.info(
-                                            "Could not record skip into DataModel "
-                                            "header: %s",
-                                            e,
-                                        )
-                            elif isinstance(args[0], AbstractDataModel):
+
+                    elif (
+                        (isinstance(args[0], Sequence))
+                        and (not isinstance(args[0], str))
+                        and (self.class_alias is not None)
+                    ):
+                        # handle ModelContainer or list of models
+                        if args[0] and isinstance(args[0][0], AbstractDataModel):
+                            for model in args[0]:
                                 try:
-                                    args[0][
-                                        f"meta.cal_step.{self.class_alias}"
-                                    ] = "SKIPPED"
+                                    setattr(
+                                        model.meta.cal_step, self.class_alias, "SKIPPED"
+                                    )
                                 except AttributeError as e:
                                     self.log.info(
-                                        "Could not record skip into DataModel"
-                                        " header: %s",
+                                        "Could not record skip into DataModel "
+                                        "header: %s",
                                         e,
                                     )
+
+                    elif (
+                        isinstance(args[0], AbstractDataModel)
+                        and self.class_alias is not None
+                    ):
+                        try:
+                            args[0][f"meta.cal_step.{self.class_alias}"] = "SKIPPED"
+                        except AttributeError as e:
+                            self.log.info(
+                                "Could not record skip into DataModel header: %s",
+                                e,
+                            )
                     step_result = args[0]
                 else:
                     if self.prefetch_references:
@@ -558,10 +564,13 @@
                 # Save the output file if one was specified
                 if not self.skip and self.save_results:
                     # Setup the save list.
-                    if not isinstance(step_result, list | tuple):
-                        results_to_save = [step_result]
+                    if isinstance(step_result, Sequence):
+                        if hasattr(step_result, "save") or isinstance(step_result, str):
+                            results_to_save = [step_result]
+                        else:
+                            results_to_save = step_result
                     else:
-                        results_to_save = step_result
+                        results_to_save = [step_result]
 
                     for idx, result in enumerate(results_to_save):
                         if len(results_to_save) <= 1:
@@ -992,6 +1001,9 @@
         if not force and not self.save_results and not output_file:
             return None
 
+        if model is None:
+            return None
+
         if isinstance(model, AbstractModelLibrary):
             output_paths = []
             with model:
@@ -1008,6 +1020,7 @@
                     # leaving modify=True in case saving modify the file
                     model.shelve(m, i)
             return output_paths
+
         elif isinstance(model, Sequence):
             save_model_func = partial(
                 self.save_model,

tests/test_step.py

@@ -1,14 +1,17 @@
 """Test step.Step"""
 
+import copy
 import logging
 import re
+from collections.abc import Sequence
 from typing import ClassVar
 
 import asdf
 import pytest
 
 import stpipe.config_parser as cp
 from stpipe import cmdline
+from stpipe.datamodel import AbstractDataModel
 from stpipe.pipeline import Pipeline
 from stpipe.step import Step
 
@@ -411,3 +414,159 @@ def test_log_records():
     pipeline.run()
 
     assert any(r == "This step has called out a warning." for r in pipeline.log_records)
+
+
+class StepWithModel(Step):
+    """A step that immediately saves the model it gets passed in"""
+
+    spec = """
+    output_ext = string(default='simplestep')
+    save_results = boolean(default=True)
+    """
+
+    def process(self, input_model):
+        # make a change to ensure step skip is working
+        # without having to define SimpleDataModel.meta.stepname
+        if isinstance(input_model, SimpleDataModel):
+            input_model.stepstatus = "COMPLETED"
+        elif isinstance(input_model, SimpleContainer):
+            for model in input_model:
+                model.stepstatus = "COMPLETED"
+        return input_model
+
+
+class SimpleDataModel(AbstractDataModel):
+    """A simple data model"""
+
+    @property
+    def crds_observatory(self):
+        return "jwst"
+
+    def get_crds_parameters(self):
+        return {"test": "none"}
+
+    def save(self, path, dir_path=None, *args, **kwargs):
+        saveid = getattr(self, "saveid", None)
+        if saveid is not None:
+            fname = saveid + "-saved.txt"
+            with open(fname, "w") as f:
+                f.write(f"{path}")
+            return fname
+        return None
+
+
+def test_save(tmp_cwd):
+
+    model = SimpleDataModel()
+    model.saveid = "test"
+    step = StepWithModel()
+    step.run(model)
+    assert (tmp_cwd / "test-saved.txt").exists()
+
+
+def test_skip():
+    model = SimpleDataModel()
+    step = StepWithModel()
+    step.skip = True
+    out = step.run(model)
+    assert not hasattr(out, "stepstatus")
+    assert out is model
+
+
+@pytest.fixture(scope="function")
+def model_list():
+    model = SimpleDataModel()
+    model_list = [copy.deepcopy(model) for _ in range(3)]
+    for i, model in enumerate(model_list):
+        model.saveid = f"test{i}"
+    return model_list
+
+
+def test_save_list(tmp_cwd, model_list):
+    step = StepWithModel()
+    step.run(model_list)
+    for i in range(3):
+        assert (tmp_cwd / f"test{i}-saved.txt").exists()
+
+
+class SimpleContainer(Sequence):
+
+    def __init__(self, models):
+        self._models = models
+
+    def __len__(self):
+        return len(self._models)
+
+    def __getitem__(self, idx):
+        return self._models[idx]
+
+    def __iter__(self):
+        yield from self._models
+
+    def insert(self, index, model):
+        self._models.insert(index, model)
+
+    def append(self, model):
+        self._models.append(model)
+
+    def extend(self, model):
+        self._models.extend(model)
+
+    def pop(self, index=-1):
+        self._models.pop(index)
+
+
+class SimpleContainerWithSave(SimpleContainer):
+
+    def save(self, path, dir_path=None, *args, **kwargs):
+        for model in self._models[1:]:
+            # skip the first model to test that the save method is called
+            # rather than just looping over all models like in the without-save case
+            model.save(path, dir_path, *args, **kwargs)
+
+
+def test_save_container(tmp_cwd, model_list):
+    """ensure list-like save still works for non-list sequence"""
+    container = SimpleContainer(model_list)
+    step = StepWithModel()
+    step.run(container)
+    for i in range(3):
+        assert (tmp_cwd / f"test{i}-saved.txt").exists()
+
+
+def test_skip_container(tmp_cwd, model_list):
+    step = StepWithModel()
+    step.skip = True
+    out = step.run(model_list)
+    assert not hasattr(out, "stepstatus")
+    for i, model in enumerate(out):
+        assert not hasattr(model, "stepstatus")
+        assert model_list[i] is model
+
+
+def test_save_container_with_save_method(tmp_cwd, model_list):
+    """ensure list-like save still works for non-list sequence"""
+    container = SimpleContainerWithSave(model_list)
+    step = StepWithModel()
+    step.run(container)
+    assert not (tmp_cwd / "test0-saved.txt").exists()
+    assert (tmp_cwd / "test1-saved.txt").exists()
+    assert (tmp_cwd / "test2-saved.txt").exists()
+
+
+def test_save_tuple_with_nested_list(tmp_cwd, model_list):
+    """
+    in rare cases, multiple outputs are returned from step as tuple.
+    One example is the jwst badpix_selfcal step, which returns one sci exposure
+    and a list containing an arbitrary number of background exposures.
+    Expected behavior in this case, at least at time of writing, is to save the
+    science exposure and ignore the list
+    """
+    single_model = SimpleDataModel()
+    single_model.saveid = "test"
+    container = (single_model, model_list)
+    step = StepWithModel()
+    step.run(container)
+    assert (tmp_cwd / "test-saved.txt").exists()
+    for i in range(3):
+        assert not (tmp_cwd / f"test{i}-saved.txt").exists()