handle Sequence input for skip and save_results

spacetelescope · Sep 11, 2024 · f4c2814 · f4c2814
1 parent 7a3fea6
commit f4c2814
Showing 1 changed file with 38 additions and 31 deletions.
diff --git a/src/stpipe/step.py b/src/stpipe/step.py
@@ -491,31 +491,34 @@ def run(self, *args):
                                         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 self.class_alias is not None:
+                        # handle ModelContainer or list of models
+                        if 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,7 +561,7 @@ def run(self, *args):
                 # 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):
+                    if not isinstance(step_result, Sequence):
                         results_to_save = [step_result]
                     else:
                         results_to_save = step_result
@@ -1008,17 +1011,21 @@ def save_model(
                     # 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,
-                suffix=suffix,
-                force=force,
-                **components,
-            )
-            output_path = model.save(
-                path=output_file,
-                save_model_func=save_model_func,
-            )
+            output_paths = []
+            for i, m in enumerate(model):
+                output_paths.append(
+                    self.save_model(
+                        m,
+                        idx=i,
+                        suffix=suffix,
+                        force=force,
+                        **components,
+                    )
+                )
+            return output_paths
+
         else:
             # Search for an output file name.
             if self.output_use_model or (