added dataclass_transform decorators to define and outputs

Author: tclose

pydra/design/python.py

@@ -1,5 +1,6 @@
 import typing as ty
 import inspect
+from typing_extensions import dataclass_transform
 import attrs
 from .base import (
     Arg,
@@ -79,6 +80,19 @@ class out(Out):
     pass
 
 
+@dataclass_transform(
+    kw_only_default=True,
+    field_specifiers=(out,),
+)
+def outputs(wrapped):
+    """Decorator to specify the output fields of a shell command is a dataclass-style type"""
+    return wrapped
+
+
+@dataclass_transform(
+    kw_only_default=True,
+    field_specifiers=(arg,),
+)
 def define(
     wrapped: type | ty.Callable | None = None,
     /,

pydra/design/shell.py

@@ -8,6 +8,7 @@
 from copy import copy
 import attrs
 import builtins
+from typing_extensions import dataclass_transform
 from fileformats.core import from_mime
 from fileformats import generic
 from fileformats.core.exceptions import FormatRecognitionError
@@ -198,6 +199,19 @@ def _validate_path_template(self, attribute, value):
             )
 
 
+@dataclass_transform(
+    kw_only_default=True,
+    field_specifiers=(out, outarg),
+)
+def outputs(wrapped):
+    """Decorator to specify the output fields of a shell command is a dataclass-style type"""
+    return wrapped
+
+
+@dataclass_transform(
+    kw_only_default=True,
+    field_specifiers=(arg,),
+)
 def define(
     wrapped: type | str | None = None,
     /,

pydra/design/tests/test_shell.py

@@ -6,7 +6,7 @@
 import cloudpickle as cp
 from pydra.design import shell
 from pydra.engine.helpers import list_fields
-from pydra.engine.specs import ShellSpec
+from pydra.engine.specs import ShellSpec, ShellOutputs
 from fileformats.generic import File, Directory, FsObject
 from fileformats import text, image
 from pydra.utils.typing import MultiInputObj
@@ -267,7 +267,8 @@ class Ls(ShellSpec["Ls.Outputs"]):
                 xor=["complete_date"],
             )
 
-            class Outputs:
+            @shell.outputs
+            class Outputs(ShellOutputs):
                 entries: list = shell.out(
                     help_string="list of entries returned by ls command",
                     callable=list_entries,
@@ -346,7 +347,14 @@ def test_shell_fields(Ls):
         ]
     )
 
-    assert [a.name for a in sorted_fields(Ls.Outputs)] == ["entries"]
+    assert [a.name for a in sorted_fields(Ls.Outputs)] == sorted(
+        [
+            "entries",
+            "stdout",
+            "stderr",
+            "return_code",
+        ]
+    )
 
 
 def test_shell_pickle_roundtrip(Ls, tmp_path):

pydra/design/tests/test_workflow.py

@@ -152,6 +152,7 @@ def constructor(a, b):
             mul = workflow.add(Mul(a=add.out, b=b))
             return mul.out
 
+        @workflow.outputs
         class Outputs(WorkflowOutputs):
             out: float
 

pydra/design/workflow.py

@@ -1,5 +1,6 @@
 import typing as ty
 import inspect
+from typing_extensions import dataclass_transform
 import attrs
 from .base import (
     Arg,
@@ -84,6 +85,19 @@ class out(Out):
     pass
 
 
+@dataclass_transform(
+    kw_only_default=True,
+    field_specifiers=(out,),
+)
+def outputs(wrapped):
+    """Decorator to specify the output fields of a shell command is a dataclass-style type"""
+    return wrapped
+
+
+@dataclass_transform(
+    kw_only_default=True,
+    field_specifiers=(arg,),
+)
 def define(
     wrapped: type | ty.Callable | None = None,
     /,