added task_def_as_dict and task_def_from_dict utility methods for serializing task classes to and from dictionaries

Author: tclose

pydra/compose/base/helpers.py

@@ -217,6 +217,8 @@ def extract_function_inputs_and_outputs(
                 inpt.default = default
         elif inspect.isclass(inpt) or ty.get_origin(inpt):
             inputs[inpt_name] = arg_type(type=inpt, default=default)
+        elif isinstance(inpt, dict):
+            inputs[inpt_name] = arg_type(**inpt)
         else:
             raise ValueError(
                 f"Unrecognised input type ({inpt}) for input {inpt_name} with default "

pydra/compose/base/task.py

@@ -149,6 +149,11 @@ def __repr__(self) -> str:
 class Task(ty.Generic[OutputsType]):
     """Base class for all tasks"""
 
+    # Task type to be overridden in derived classes
+    _task_type = ""
+    # The attribute containing the function/executable used to run the task
+    _executor_name = None
+
     # Class attributes
     _xor: frozenset[frozenset[str | None]] = (
         frozenset()

pydra/compose/python.py

@@ -102,6 +102,7 @@ def define(
     bases: ty.Sequence[type] = (),
     outputs_bases: ty.Sequence[type] = (),
     auto_attribs: bool = True,
+    name: str | None = None,
     xor: ty.Sequence[str | None] | ty.Sequence[ty.Sequence[str | None]] = (),
 ) -> "Task":
     """
@@ -117,6 +118,8 @@ def define(
         The outputs of the function or class.
     auto_attribs : bool
         Whether to use auto_attribs mode when creating the class.
+    name: str | None
+        The name of the returned class
     xor: Sequence[str | None] | Sequence[Sequence[str | None]], optional
         Names of args that are exclusive mutually exclusive, which must include
         the name of the current field. If this list includes None, then none of the
@@ -132,7 +135,7 @@ def make(wrapped: ty.Callable | type) -> Task:
         if inspect.isclass(wrapped):
             klass = wrapped
             function = klass.function
-            name = klass.__name__
+            class_name = klass.__name__
             check_explicit_fields_are_none(klass, inputs, outputs)
             parsed_inputs, parsed_outputs = extract_fields_from_class(
                 Task,
@@ -154,7 +157,8 @@ def make(wrapped: ty.Callable | type) -> Task:
             inferred_inputs, inferred_outputs = extract_function_inputs_and_outputs(
                 function, arg, inputs, outputs
             )
-            name = function.__name__
+
+            class_name = function.__name__ if name is None else name
 
             parsed_inputs, parsed_outputs = ensure_field_objects(
                 arg_type=arg,
@@ -179,7 +183,7 @@ def make(wrapped: ty.Callable | type) -> Task:
             Outputs,
             parsed_inputs,
             parsed_outputs,
-            name=name,
+            name=class_name,
             klass=klass,
             bases=bases,
             outputs_bases=outputs_bases,
@@ -228,6 +232,7 @@ def _from_task(cls, job: "Job[PythonTask]") -> ty.Self:
 class PythonTask(base.Task[PythonOutputsType]):
 
     _task_type = "python"
+    _executor_name = "function"
 
     def _run(self, job: "Job[PythonTask]", rerun: bool = True) -> None:
         # Prepare the inputs to the function

pydra/compose/shell/task.py

@@ -233,6 +233,7 @@ def append_args_converter(value: ty.Any) -> list[str]:
 class ShellTask(base.Task[ShellOutputsType]):
 
     _task_type = "shell"
+    _executor_name = "executable"
 
     BASE_NAMES = ["append_args"]
 

pydra/compose/workflow.py

@@ -111,6 +111,7 @@ def define(
     outputs_bases: ty.Sequence[type] = (),
     lazy: list[str] | None = None,
     auto_attribs: bool = True,
+    name: str | None = None,
     xor: ty.Sequence[str | None] | ty.Sequence[ty.Sequence[str | None]] = (),
 ) -> "Task":
     """
@@ -127,6 +128,8 @@ def define(
         The outputs of the function or class.
     auto_attribs : bool
         Whether to use auto_attribs mode when creating the class.
+    name: str | None
+        The name of the returned class
     xor: Sequence[str | None] | Sequence[Sequence[str | None]], optional
         Names of args that are exclusive mutually exclusive, which must include
         the name of the current field. If this list includes None, then none of the
@@ -145,7 +148,7 @@ def make(wrapped: ty.Callable | type) -> Task:
         if inspect.isclass(wrapped):
             klass = wrapped
             constructor = klass.constructor
-            name = klass.__name__
+            class_name = klass.__name__
             check_explicit_fields_are_none(klass, inputs, outputs)
             parsed_inputs, parsed_outputs = extract_fields_from_class(
                 Task,
@@ -167,7 +170,8 @@ def make(wrapped: ty.Callable | type) -> Task:
             inferred_inputs, inferred_outputs = extract_function_inputs_and_outputs(
                 constructor, arg, inputs, outputs
             )
-            name = constructor.__name__
+
+            class_name = constructor.__name__ if name is None else name
 
             parsed_inputs, parsed_outputs = ensure_field_objects(
                 arg_type=arg,
@@ -195,7 +199,7 @@ def make(wrapped: ty.Callable | type) -> Task:
             Outputs,
             parsed_inputs,
             parsed_outputs,
-            name=name,
+            name=class_name,
             klass=klass,
             bases=bases,
             outputs_bases=outputs_bases,
@@ -349,6 +353,7 @@ def _from_task(cls, job: "Job[WorkflowTask]") -> ty.Self:
 class WorkflowTask(base.Task[WorkflowOutputsType]):
 
     _task_type = "workflow"
+    _executor_name = "constructor"
 
     RESERVED_FIELD_NAMES = base.Task.RESERVED_FIELD_NAMES + ("construct",)
 

pydra/utils/general.py

@@ -4,6 +4,7 @@
 import inspect
 import sys
 import typing as ty
+from copy import copy
 import re
 import attrs
 import ast
@@ -584,3 +585,61 @@ def get_plugin_classes(namespace: types.ModuleType, class_name: str) -> dict[str
         for pkg in sub_packages
         if hasattr(pkg, class_name)
     }
+
+
+def task_def_as_dict(task_def: "type[Task]") -> ty.Dict[str, ty.Any]:
+    """Converts a Pydra task class into a dictionary representation that can be serialized
+    and saved to a file, then read and passed to an appropriate `pydra.compose.*.define`
+    method to recreate the task.
+
+    Parameters
+    ----------
+    task_def : type[pydra.compose.base.Task]
+        The Pydra task class to convert.
+
+    Returns
+    -------
+    dict[str, ty.Any]
+        A dictionary representation of the Pydra task.
+    """
+    input_fields = task_fields(task_def)
+    executor = input_fields.pop(task_def._executor_name).default
+    input_dicts = [attrs.asdict(i, filter=_filter_defaults) for i in input_fields]
+    output_dicts = [
+        attrs.asdict(o, filter=_filter_defaults) for o in task_fields(task_def.Outputs)
+    ]
+    dct = {
+        "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,
+    }
+
+    return dct
+
+
+def task_def_from_dict(task_def_dict: dict[str, ty.Any]) -> type["Task"]:
+    """Unserializes a task definition from a dictionary created by `task_def_as_dict`
+
+    Parameters
+    ----------
+    task_def_dict: dict[str, Any]
+        the dictionary representation to unserialize
+
+    Returns
+    -------
+    type[pydra.compose.base.Task]
+        the unserialized task class
+    """
+    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("function"), **dct)
+
+
+def _filter_defaults(atr: attrs.Attribute, value: ty.Any) -> bool:
+    if value == atr.default:
+        return False
+    return True

pydra/utils/tests/test_general.py

@@ -0,0 +1,28 @@
+from pydra.compose import python  # , workflow, shell
+from pydra.utils.general import task_def_as_dict, task_def_from_dict, task_fields
+
+
+def test_python_task_def_as_dict():
+
+    @python.define(outputs=["out_int"], xor=["b", "c"])
+    def Add(a: int, b: int | None = None, c: int | None = None) -> int:
+        """
+        Parameters
+        ----------
+        a: int
+            the first arg
+        b : int, optional
+            the optional second arg
+        c : int, optional
+            the optional third arg
+
+        Returns
+        -------
+        out_int : int
+            the sum of a and b
+        """
+        return a + (b if b is not None else c)
+
+    dct = task_def_as_dict(Add)
+    Reloaded = task_def_from_dict(dct)
+    assert task_fields(Add) == task_fields(Reloaded)