From f4c281442fd02ebbf5822b8a2ec758bab1634560 Mon Sep 17 00:00:00 2001 From: Ned Molter Date: Wed, 11 Sep 2024 10:54:11 -0400 Subject: [PATCH] handle Sequence input for skip and save_results --- src/stpipe/step.py | 69 +++++++++++++++++++++++++--------------------- 1 file changed, 38 insertions(+), 31 deletions(-) diff --git a/src/stpipe/step.py b/src/stpipe/step.py index f956c454..7f0d508a 100644 --- 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 (