Merge pull request #792 from nipype/task-def-dict-util

Serializing Task classes to/from dictionaries

Author: tclose

README.rst

@@ -106,6 +106,7 @@ It is also useful to install pre-commit:
 
 .. |codecov| image:: https://codecov.io/gh/nipype/pydra/branch/main/graph/badge.svg
    :alt: codecov
+   :target: https://codecov.io/gh/nipype/pydra
 
 .. |Python Versions| image:: https://img.shields.io/pypi/pyversions/pydra.svg
    :alt: Supported Python versions

docs/source/index.rst

@@ -85,7 +85,7 @@ workflows,
 * :ref:`Python-tasks`
 * :ref:`Shell-tasks`
 * :ref:`Workflows`
-* :ref:`Canonical task form`
+* :ref:`Canonical form and serialisation`
 
 Examples
 ~~~~~~~~

docs/source/tutorial/5-shell.ipynb

@@ -162,7 +162,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from pydra.utils import task_fields\n",
+    "from pydra.utils import get_fields\n",
     "\n",
     "Cp = shell.define(\n",
     "    \"cp <in_fs_objects:fs-object+> <out|out_dir:directory> \"\n",
@@ -172,7 +172,7 @@
     "    \"--tuple-arg <tuple_arg:int,str=(1,'bar')> \"\n",
     ")\n",
     "\n",
-    "print(f\"'--int-arg' default: {task_fields(Cp).int_arg.default}\")"
+    "print(f\"'--int-arg' default: {get_fields(Cp).int_arg.default}\")"
    ]
   },
   {

docs/source/tutorial/7-canonical-form.ipynb

@@ -4,7 +4,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Canonical task form\n",
+    "# Canonical form and serialisation\n",
+    "\n",
+    "## Canonical task form\n",
     "\n",
     "Under the hood, all Python, shell and workflow tasks generated by the\n",
     "`pydra.compose.*.define` decorators/functions are translated to\n",
@@ -149,33 +151,32 @@
     "from pathlib import Path\n",
     "from fileformats import generic\n",
     "from pydra.compose import shell\n",
-    "from pydra.utils.typing import MultiInputObj\n",
+    "\n",
+    "\n",
+    "# the \"copied\" output is magically passed to this function because the name matches\n",
+    "def get_file_size(copied: Path) -> int:\n",
+    "    \"\"\"Calculate the file size\"\"\"\n",
+    "    result = os.stat(copied)\n",
+    "    return result.st_size\n",
     "\n",
     "\n",
     "@shell.define\n",
-    "class CpWithSize(shell.Task[\"CpWithSize.Outputs\"]):\n",
+    "class CpFileWithSize(shell.Task[\"CpFileWithSize.Outputs\"]):\n",
     "\n",
     "    executable = \"cp\"\n",
     "\n",
-    "    in_fs_objects: MultiInputObj[generic.FsObject]\n",
-    "    recursive: bool = shell.arg(argstr=\"-R\")\n",
-    "    text_arg: str = shell.arg(argstr=\"--text-arg\")\n",
-    "    int_arg: int | None = shell.arg(argstr=\"--int-arg\")\n",
-    "    tuple_arg: tuple[int, str] | None = shell.arg(argstr=\"--tuple-arg\")\n",
+    "    in_file: generic.File  # = shell.arg() is assumed\n",
+    "    archive_mode: bool = shell.arg(argstr=\"-a\", default=False)\n",
     "\n",
     "    class Outputs(shell.Outputs):\n",
     "\n",
-    "        @staticmethod\n",
-    "        def get_file_size(out_file: Path) -> int:\n",
-    "            \"\"\"Calculate the file size\"\"\"\n",
-    "            result = os.stat(out_file)\n",
-    "            return result.st_size\n",
+    "        copied: generic.File = shell.outarg(\n",
+    "            position=-1, path_template=\"{in_file}_copied\"\n",
+    "        )\n",
+    "        file_size: int = shell.out(callable=get_file_size)\n",
     "\n",
-    "        copied: generic.FsObject = shell.outarg(path_template=\"copied\")\n",
-    "        out_file_size: int = shell.out(callable=get_file_size)\n",
     "\n",
-    "\n",
-    "print_help(CpWithSize)"
+    "print_help(CpFileWithSize)"
    ]
   },
   {
@@ -197,7 +198,6 @@
     "import typing as ty\n",
     "import re\n",
     "from pydra.compose import python, workflow\n",
-    "from pydra.compose.base import is_set\n",
     "from pydra.utils import print_help, show_workflow\n",
     "\n",
     "\n",
@@ -237,6 +237,120 @@
     "print_help(CanonicalWorkflowTask)\n",
     "show_workflow(CanonicalWorkflowTask)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Serialization\n",
+    "\n",
+    "As well as the dataclass-like canonical form, it is also possible to represent all tasks\n",
+    "in a nested dictionary form, which could be written to a static file (e.g. in JSON or\n",
+    "YAML format). The dictionary form of a class can be generated by the `pydra.utils.unstructure`\n",
+    "function. For example, the following shell command"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MyCmd = shell.define(\n",
+    "    \"my-cmd <in_file> <out|out_file> --an-arg <an_arg?> \"\n",
+    "    \"--a-flag<a_flag> --arg-with-default <arg_with_default:int=3>\"\n",
+    ")\n",
+    "\n",
+    "print_help(MyCmd)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Can be converted into a serialised dictionary form"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pprint import pprint\n",
+    "from pydra.utils import unstructure\n",
+    "\n",
+    "my_cmd_dict = unstructure(MyCmd)\n",
+    "\n",
+    "pprint(my_cmd_dict)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Noting that there is still a little more work has to be done to serialise some Python\n",
+    "objects, e.g. classes used in field types and functions that are run in Python and\n",
+    "construct workflows in workflow tasks, before the serialized form can be written to JSON/YAML."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cp_with_size_dict = unstructure(CpFileWithSize)\n",
+    "\n",
+    "pprint(cp_with_size_dict)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To unserialize the general dictionary form back into a Task class, you can use the\n",
+    "`pydra.utils.structure` method"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pydra.utils import structure\n",
+    "\n",
+    "ReloadedCpFileWithSize = structure(cp_with_size_dict)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "which should run just as before"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "import tempfile\n",
+    "from pydra.utils import asdict\n",
+    "\n",
+    "tmp_dir = Path(tempfile.mkdtemp())\n",
+    "\n",
+    "a_file = tmp_dir / \"hello-world.txt\"\n",
+    "a_file.write_text(\"Hello world\")\n",
+    "\n",
+    "cp_file_with_size = ReloadedCpFileWithSize(in_file=a_file)\n",
+    "outputs = cp_file_with_size(cache_root=tmp_dir / \"cache\")\n",
+    "\n",
+    "pprint(asdict(outputs))"
+   ]
   }
  ],
  "metadata": {

pydra/compose/base/field.py

@@ -6,7 +6,7 @@
 from fileformats.core import to_mime
 from fileformats.generic import File, FileSet
 from pydra.utils.typing import TypeParser, is_optional, is_type, is_union
-from pydra.utils.general import task_fields, wrap_text
+from pydra.utils.general import get_fields, wrap_text
 import attrs
 
 if ty.TYPE_CHECKING:
@@ -66,7 +66,7 @@ class Requirement:
     def satisfied(self, inputs: "Task") -> bool:
         """Check if the requirement is satisfied by the inputs"""
         value = getattr(inputs, self.name)
-        field = {f.name: f for f in task_fields(inputs)}[self.name]
+        field = {f.name: f for f in get_fields(inputs)}[self.name]
         if value is None or field.type is bool and value is False:
             return False
         if self.allowed_values is None:
@@ -326,7 +326,7 @@ class Arg(Field):
         it is False
     """
 
-    allowed_values: frozenset = attrs.field(default=(), converter=frozenset)
+    allowed_values: frozenset = attrs.field(factory=frozenset, converter=frozenset)
     copy_mode: File.CopyMode = File.CopyMode.any
     copy_collation: File.CopyCollation = File.CopyCollation.any
     copy_ext_decomp: File.ExtensionDecomposition = File.ExtensionDecomposition.single

pydra/compose/base/helpers.py

@@ -4,7 +4,7 @@
 import re
 from copy import copy
 from pydra.utils.typing import is_type, is_optional
-from pydra.utils.general import task_fields
+from pydra.utils.general import get_fields
 from .field import Field, Arg, Out, NO_DEFAULT
 
 
@@ -114,7 +114,13 @@ def ensure_field_objects(
             out_kwds = copy(out)
             if "help" not in out_kwds:
                 out_kwds["help"] = output_helps.get(output_name, "")
-            outputs[output_name] = out_type(
+            if "path_template" in out_kwds:
+                from pydra.compose.shell import outarg
+
+                out_type_ = outarg
+            else:
+                out_type_ = out_type
+            outputs[output_name] = out_type_(
                 name=output_name,
                 **out_kwds,
             )
@@ -217,6 +223,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 "
@@ -406,11 +414,11 @@ def extract_fields_from_class(
 
     input_helps, _ = parse_doc_string(klass.__doc__)
 
-    def get_fields(klass, field_type, auto_attribs, helps) -> dict[str, Field]:
+    def extract_fields(klass, field_type, auto_attribs, helps) -> dict[str, Field]:
         """Get the fields from a class"""
         fields_dict = {}
         # Get fields defined in base classes if present
-        for field in task_fields(klass):
+        for field in get_fields(klass):
             if field.name not in skip_fields:
                 fields_dict[field.name] = field
         type_hints = ty.get_type_hints(klass)
@@ -460,7 +468,7 @@ def get_fields(klass, field_type, auto_attribs, helps) -> dict[str, Field]:
             f"tasks, {klass} must inherit from {spec_type}"
         )
 
-    inputs = get_fields(klass, arg_type, auto_attribs, input_helps)
+    inputs = extract_fields(klass, arg_type, auto_attribs, input_helps)
 
     try:
         outputs_klass = klass.Outputs
@@ -475,7 +483,7 @@ def get_fields(klass, field_type, auto_attribs, helps) -> dict[str, Field]:
         )
 
     output_helps, _ = parse_doc_string(outputs_klass.__doc__)
-    outputs = get_fields(outputs_klass, out_type, auto_attribs, output_helps)
+    outputs = extract_fields(outputs_klass, out_type, auto_attribs, output_helps)
 
     return inputs, outputs