standardize setattr use in skip

spacetelescope · Oct 9, 2024 · 7a5ac2a · 7a5ac2a
1 parent ecd5d16
commit 7a5ac2a
Showing 1 changed file with 16 additions and 39 deletions.
diff --git a/src/stpipe/step.py b/src/stpipe/step.py
@@ -474,48 +474,25 @@ def run(self, *args):
                 self._check_args(args, DISCOURAGED_TYPES, "Passed")
 
                 # Run the Step-specific code.
-                if self.skip:
+                if self.skip and self.class_alias is not None:
                     self.log.info("Step skipped.")
+
+                    def set_skipped(model):
+                        try:
+                            setattr(model.meta.cal_step, self.class_alias, "SKIPPED")
+                        except AttributeError as e:
+                            self.log.info(
+                                "Could not record skip into DataModel " "header: %s",
+                                e,
+                            )
+
                     if isinstance(args[0], AbstractModelLibrary):
-                        library = args[0]
-                        with library:
-                            for i, model in enumerate(library):
-                                try:
-                                    setattr(
-                                        model.meta.cal_step, self.class_alias, "SKIPPED"
-                                    )
-                                except AttributeError as e:
-                                    self.log.info(
-                                        "Could not record skip into DataModel "
-                                        "header: %s",
-                                        e,
-                                    )
-                                library.shelve(model, i)
+                        list(args[0].map_function(lambda m, i: set_skipped(m)))
                     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):
-                                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,
-                                    )
+                        if isinstance(args[0], Sequence):
+                            [set_skipped(m) for m in args[0]]
+                        else:
+                            set_skipped(args[0])
                     step_result = args[0]
                 else:
                     if self.prefetch_references: