fixed broken _task_type and _from_job

Author: tclose

pydra/compose/base/task.py

@@ -152,9 +152,15 @@ class Task(ty.Generic[OutputsType]):
     # Task type to be overridden in derived classes
     @classmethod
     def _task_type(cls) -> str:
-        mod_parts = cls.__module__.split(".")
-        assert mod_parts[:2] == ["pydra", "compose"]
-        return mod_parts[2]
+        for base in cls.__mro__:
+            parts = base.__module__.split(".")
+            if parts[:2] == ["pydra", "compose"]:
+                return parts[2]
+        raise RuntimeError(
+            f"Cannot determine task type for {cls.__name__} in module {cls.__module__} "
+            "because none of its base classes are in the pydra.compose namespace:\n"
+            + "\n".join(f"{b.__name__!r} in {b.__module__!r}" for b in cls.__mro__)
+        )
 
     # The attribute containing the function/executable used to run the task
     _executor_name = None

pydra/compose/python.py

@@ -219,7 +219,7 @@ def _from_job(cls, job: "Job[PythonTask]") -> ty.Self:
         outputs : Outputs
             The outputs of the job in dataclass
         """
-        outputs = super()._from_task(job)
+        outputs = super()._from_job(job)
         for name, val in job.return_values.items():
             setattr(outputs, name, val)
         return outputs

pydra/compose/shell/task.py

@@ -78,7 +78,7 @@ def _from_job(cls, job: "Job[Task]") -> ty.Self:
         outputs : Outputs
             The outputs of the shell process
         """
-        outputs = super()._from_task(job)
+        outputs = super()._from_job(job)
         fld: field.out
         for fld in task_fields(cls):
             if fld.name in ["return_code", "stdout", "stderr"]:

pydra/compose/workflow.py

@@ -340,7 +340,7 @@ def _from_job(cls, job: "Job[WorkflowTask]") -> ty.Self:
             values[name] = val_out
 
         # Set the values in the outputs object
-        outputs = super()._from_task(job)
+        outputs = super()._from_job(job)
         outputs = attrs.evolve(outputs, **values)
         outputs._cache_dir = job.cache_dir
         return outputs

pydra/engine/job.py

@@ -331,7 +331,7 @@ def run(self, rerun: bool = False):
             try:
                 self.audit.monitor()
                 self.task._run(self, rerun)
-                result.outputs = self.task.Outputs._from_task(self)
+                result.outputs = self.task.Outputs._from_job(self)
             except Exception:
                 etype, eval, etr = sys.exc_info()
                 traceback = format_exception(etype, eval, etr)
@@ -385,7 +385,7 @@ async def run_async(self, rerun: bool = False) -> Result:
             try:
                 self.audit.monitor()
                 await self.task._run_async(self, rerun)
-                result.outputs = self.task.Outputs._from_task(self)
+                result.outputs = self.task.Outputs._from_job(self)
             except Exception:
                 etype, eval, etr = sys.exc_info()
                 traceback = format_exception(etype, eval, etr)