apply value_serializer to class attributes

Author: tclose

pydra/utils/general.py

@@ -650,9 +650,15 @@ def task_class_as_dict(
     }
     class_attrs = {a: getattr(task_class, "_" + a) for a in task_class.TASK_CLASS_ATTRS}
     if value_serializer:
-        attrs_fields = {f.name: f for f in attrs.fields(task_class)}
+        # We need to create a mock attrs object for the class attrs to apply the
+        # value_serializer to it
+        mock_cls = _make_attrs_class(
+            {a: task_class.__annotations__.get(a, ty.Any) for a in class_attrs}
+        )
+        attrs_fields = {f.name: f for f in attrs.fields(mock_cls)}
+        mock = mock_cls(**class_attrs)
         class_attrs = {
-            n: value_serializer(task_class, attrs_fields[n], v)
+            n: value_serializer(mock, attrs_fields[n], v)
             for n, v in class_attrs.items()
         }
     dct.update(class_attrs)
@@ -686,3 +692,25 @@ def _filter_out_defaults(atr: attrs.Attribute, value: ty.Any) -> bool:
     if value == atr.default:
         return False
     return True
+
+
+def _make_attrs_class(field_types: dict[str, type]) -> type:
+    """Creates an attrs given the a dictionary of field names and their types.
+
+    Parameters
+    ----------
+    field_types : dict[str, type]
+        A dictionary mapping field names to their types.
+
+    Returns
+    -------
+    type
+        the attrs class.
+    """
+    return attrs.define(
+        type(
+            "MockAttrsClass",
+            (),
+            {n: attrs.field(type=t) for n, t in field_types.items()},
+        )
+    )

pydra/utils/tests/test_general.py

@@ -1,34 +1,40 @@
+import typing as ty
+import attrs
 from pydra.compose import python, workflow, shell
-from fileformats.generic import File
+from fileformats.text import TextFile
 from pydra.utils.general import task_class_as_dict, task_class_from_dict, task_fields
-from pydra.utils.tests.utils import SpecificFuncTask, Concatenate
-
-
-def test_python_task_class_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)
+from pydra.utils.tests.utils import Concatenate
+
+
+@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)
+
+
+def test_python_task_class_as_dict(tmp_path):
 
     dct = task_class_as_dict(Add)
     Reloaded = task_class_from_dict(dct)
     assert task_fields(Add) == task_fields(Reloaded)
 
+    add = Reloaded(a=1, b=2)
+    assert add(cache_root=tmp_path / "cache").out_int == 3
+
 
 def test_shell_task_class_as_dict():
 
@@ -41,18 +47,37 @@ def test_shell_task_class_as_dict():
     assert task_fields(MyCmd) == task_fields(Reloaded)
 
 
-def test_workflow_task_class_as_dict():
+def test_workflow_task_class_as_dict(tmp_path):
 
-    @workflow.define
-    def AWorkflow(in_file: File, a_param: int) -> tuple[File, File]:
-        spec_func = workflow.add(SpecificFuncTask(in_file))
+    @workflow.define(outputs=["out_file"])
+    def AWorkflow(in_file: TextFile, a_param: int) -> TextFile:
         concatenate = workflow.add(
-            Concatenate(
-                in_file1=in_file, in_file2=spec_func.out_file, duplicates=a_param
-            )
+            Concatenate(in_file1=in_file, in_file2=in_file, duplicates=a_param)
         )
         return concatenate.out_file
 
     dct = task_class_as_dict(AWorkflow)
     Reloaded = task_class_from_dict(dct)
     assert task_fields(AWorkflow) == task_fields(Reloaded)
+
+    foo_file = tmp_path / "file1.txt"
+    foo_file.write_text("foo")
+
+    outputs = Reloaded(in_file=foo_file, a_param=2)(cache_root=tmp_path / "cache")
+    assert outputs.out_file.contents == "foo\nfoo\nfoo\nfoo"
+
+
+def test_task_class_as_dict_with_value_serializer():
+
+    def frozen_set_to_list_serializer(
+        mock_class: ty.Any, atr: attrs.Attribute, value: ty.Any
+    ) -> ty.Any:
+        # This is just a dummy serializer
+        if isinstance(value, frozenset):
+            return list(
+                frozen_set_to_list_serializer(mock_class, atr, v) for v in value
+            )
+        return value
+
+    dct = task_class_as_dict(Add, value_serializer=frozen_set_to_list_serializer)
+    assert dct["xor"] == [["b", "c"]] or dct["xor"] == [["c", "b"]]