automated _task_type so that it always matches the name of the sub-package of pydra.compose the Task is defined in

Author: tclose

pydra/compose/base/task.py

@@ -150,11 +150,18 @@ class Task(ty.Generic[OutputsType]):
     """Base class for all tasks"""
 
     # Task type to be overridden in derived classes
-    _task_type = ""
+    @classmethod
+    def _task_type(cls) -> str:
+        mod_parts = cls.__module__.split(".")
+        assert len(mod_parts) == 3
+        assert mod_parts[:2] == ["pydra", "compose"]
+        return mod_parts[2]
+
     # The attribute containing the function/executable used to run the task
     _executor_name = None
 
     # Class attributes
+    TASK_CLASS_ATTRS = ("xor",)
     _xor: frozenset[frozenset[str | None]] = (
         frozenset()
     )  # overwritten in derived classes
@@ -461,7 +468,7 @@ def _hash(self):
 
     @property
     def _checksum(self):
-        return f"{self._task_type}-{self._hash}"
+        return f"{self._task_type()}-{self._hash}"
 
     def _hash_changes(self):
         """Detects any changes in the hashed values between the current inputs and the
@@ -610,7 +617,7 @@ def _check_resolved(self):
 
 @register_serializer
 def bytes_repr_task(obj: Task, cache: Cache) -> ty.Iterator[bytes]:
-    yield f"task[{obj._task_type}]:(".encode()
+    yield f"task[{obj._task_type()}]:(".encode()
     for field in task_fields(obj):
         yield f"{field.name}=".encode()
         yield hash_single(getattr(obj, field.name), cache)

pydra/compose/python.py

@@ -231,7 +231,6 @@ def _from_task(cls, job: "Job[PythonTask]") -> ty.Self:
 @attrs.define(kw_only=True, auto_attribs=False, eq=False, repr=False)
 class PythonTask(base.Task[PythonOutputsType]):
 
-    _task_type = "python"
     _executor_name = "function"
 
     def _run(self, job: "Job[PythonTask]", rerun: bool = True) -> None:

pydra/compose/shell/task.py

@@ -232,7 +232,6 @@ def append_args_converter(value: ty.Any) -> list[str]:
 @attrs.define(kw_only=True, auto_attribs=False, eq=False, repr=False)
 class ShellTask(base.Task[ShellOutputsType]):
 
-    _task_type = "shell"
     _executor_name = "executable"
 
     BASE_NAMES = ["append_args"]

pydra/compose/workflow.py

@@ -352,7 +352,6 @@ def _from_task(cls, job: "Job[WorkflowTask]") -> ty.Self:
 @attrs.define(kw_only=True, auto_attribs=False, eq=False, repr=False)
 class WorkflowTask(base.Task[WorkflowOutputsType]):
 
-    _task_type = "workflow"
     _executor_name = "constructor"
 
     RESERVED_FIELD_NAMES = base.Task.RESERVED_FIELD_NAMES + ("construct",)

pydra/engine/workflow.py

@@ -251,11 +251,11 @@ def add(
         if (
             environment
             and not isinstance(environment, native.Environment)
-            and task._task_type != "shell"
+            and task._task_type() != "shell"
         ):
             raise ValueError(
                 "Environments can only be used with 'shell' tasks not "
-                f"{task._task_type!r} tasks ({task!r})"
+                f"{task._task_type()!r} tasks ({task!r})"
             )
         node = Node[OutputsType](
             name=name,

pydra/utils/general.py

@@ -609,13 +609,13 @@ def task_def_as_dict(task_def: "type[Task]") -> ty.Dict[str, ty.Any]:
         attrs.asdict(o, filter=_filter_defaults) for o in task_fields(task_def.Outputs)
     ]
     dct = {
-        "type": task_def._task_type,
+        "type": task_def._task_type(),
         task_def._executor_name: executor,
         "name": task_def.__name__,
         "inputs": {d.pop("name"): d for d in input_dicts},
         "outputs": {d.pop("name"): d for d in output_dicts},
-        "xor": task_def._xor,
     }
+    dct.update({a: getattr(task_def, "_" + a) for a in task_def.TASK_CLASS_ATTRS})
 
     return dct
 
@@ -635,8 +635,8 @@ def task_def_from_dict(task_def_dict: dict[str, ty.Any]) -> type["Task"]:
     """
     dct = copy(task_def_dict)
     task_type = dct.pop("type")
-    compose_module = importlib.import_module(f"pydra.compose.{task_type}")
-    return compose_module.define(dct.pop(compose_module.Task._executor_name), **dct)
+    mod = importlib.import_module(f"pydra.compose.{task_type}")
+    return mod.define(dct.pop(mod.Task._executor_name), **dct)
 
 
 def _filter_defaults(atr: attrs.Attribute, value: ty.Any) -> bool: