Skip to content

Commit f227987

Browse files
committed
added dataclass_transform decorators to define and outputs
1 parent 51c911b commit f227987

File tree

5 files changed

+54
-3
lines changed

5 files changed

+54
-3
lines changed

pydra/design/python.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import typing as ty
22
import inspect
3+
from typing_extensions import dataclass_transform
34
import attrs
45
from .base import (
56
Arg,
@@ -79,6 +80,19 @@ class out(Out):
7980
pass
8081

8182

83+
@dataclass_transform(
84+
kw_only_default=True,
85+
field_specifiers=(out,),
86+
)
87+
def outputs(wrapped):
88+
"""Decorator to specify the output fields of a shell command is a dataclass-style type"""
89+
return wrapped
90+
91+
92+
@dataclass_transform(
93+
kw_only_default=True,
94+
field_specifiers=(arg,),
95+
)
8296
def define(
8397
wrapped: type | ty.Callable | None = None,
8498
/,

pydra/design/shell.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from copy import copy
99
import attrs
1010
import builtins
11+
from typing_extensions import dataclass_transform
1112
from fileformats.core import from_mime
1213
from fileformats import generic
1314
from fileformats.core.exceptions import FormatRecognitionError
@@ -198,6 +199,19 @@ def _validate_path_template(self, attribute, value):
198199
)
199200

200201

202+
@dataclass_transform(
203+
kw_only_default=True,
204+
field_specifiers=(out, outarg),
205+
)
206+
def outputs(wrapped):
207+
"""Decorator to specify the output fields of a shell command is a dataclass-style type"""
208+
return wrapped
209+
210+
211+
@dataclass_transform(
212+
kw_only_default=True,
213+
field_specifiers=(arg,),
214+
)
201215
def define(
202216
wrapped: type | str | None = None,
203217
/,

pydra/design/tests/test_shell.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import cloudpickle as cp
77
from pydra.design import shell
88
from pydra.engine.helpers import list_fields
9-
from pydra.engine.specs import ShellSpec
9+
from pydra.engine.specs import ShellSpec, ShellOutputs
1010
from fileformats.generic import File, Directory, FsObject
1111
from fileformats import text, image
1212
from pydra.utils.typing import MultiInputObj
@@ -267,7 +267,8 @@ class Ls(ShellSpec["Ls.Outputs"]):
267267
xor=["complete_date"],
268268
)
269269

270-
class Outputs:
270+
@shell.outputs
271+
class Outputs(ShellOutputs):
271272
entries: list = shell.out(
272273
help_string="list of entries returned by ls command",
273274
callable=list_entries,
@@ -346,7 +347,14 @@ def test_shell_fields(Ls):
346347
]
347348
)
348349

349-
assert [a.name for a in sorted_fields(Ls.Outputs)] == ["entries"]
350+
assert [a.name for a in sorted_fields(Ls.Outputs)] == sorted(
351+
[
352+
"entries",
353+
"stdout",
354+
"stderr",
355+
"return_code",
356+
]
357+
)
350358

351359

352360
def test_shell_pickle_roundtrip(Ls, tmp_path):

pydra/design/tests/test_workflow.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,7 @@ def constructor(a, b):
152152
mul = workflow.add(Mul(a=add.out, b=b))
153153
return mul.out
154154

155+
@workflow.outputs
155156
class Outputs(WorkflowOutputs):
156157
out: float
157158

pydra/design/workflow.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import typing as ty
22
import inspect
3+
from typing_extensions import dataclass_transform
34
import attrs
45
from .base import (
56
Arg,
@@ -84,6 +85,19 @@ class out(Out):
8485
pass
8586

8687

88+
@dataclass_transform(
89+
kw_only_default=True,
90+
field_specifiers=(out,),
91+
)
92+
def outputs(wrapped):
93+
"""Decorator to specify the output fields of a shell command is a dataclass-style type"""
94+
return wrapped
95+
96+
97+
@dataclass_transform(
98+
kw_only_default=True,
99+
field_specifiers=(arg,),
100+
)
87101
def define(
88102
wrapped: type | ty.Callable | None = None,
89103
/,

0 commit comments

Comments
 (0)