Skip to content

Commit

Permalink
added lazy inputs test
Browse files Browse the repository at this point in the history
  • Loading branch information
tclose committed Nov 19, 2024
1 parent ac08c18 commit a77be6a
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 5 deletions.
54 changes: 53 additions & 1 deletion pydra/design/tests/test_workflow.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from operator import attrgetter
import pytest
import attrs
from pydra.engine.workflow import Workflow
from pydra.engine.workflow import Workflow, WORKFLOW_LZIN
from pydra.engine.specs import LazyField
import typing as ty
from pydra.design import shell, python, workflow, list_fields, TaskSpec
Expand Down Expand Up @@ -115,6 +115,7 @@ def MyTestShellWorkflow(


def test_workflow_canonical():
"""Test class-based workflow definition"""

# NB: We use PascalCase (i.e. class names) as it is translated into a class

Expand Down Expand Up @@ -177,6 +178,57 @@ class Outputs:
assert list(wf.node_names) == ["Add", "Mul"]


def test_workflow_lazy():

@workflow.define(lazy=["input_video", "watermark"])
def MyTestShellWorkflow(
input_video: video.Mp4,
watermark: image.Png,
watermark_dims: tuple[int, int] = (10, 10),
) -> video.Mp4:

add_watermark = workflow.add(
shell.define(
"ffmpeg -i <in_video> -i <watermark:image/png> "
"-filter_complex <filter> <out|out_video>"
)(
in_video=input_video,
watermark=watermark,
filter="overlay={}:{}".format(*watermark_dims),
),
name="add_watermark",
)
output_video = workflow.add(
shell.define(
"HandBrakeCLI -i <in_video> -o <out|out_video> "
"--width <width:int> --height <height:int>",
# By default any input/output specified with a flag (e.g. -i <in_video>)
# is considered optional, i.e. of type `FsObject | None`, and therefore
# won't be used by default. By overriding this with non-optional types,
# the fields are specified as being required.
inputs={"in_video": video.Mp4},
outputs={"out_video": video.Mp4},
)(in_video=add_watermark.out_video, width=1280, height=720),
name="resize",
).out_video

return output_video # test implicit detection of output name

input_video = video.Mp4.mock("input.mp4")
watermark = image.Png.mock("watermark.png")
workflow_spec = MyTestShellWorkflow(
input_video=input_video,
watermark=watermark,
)
wf = Workflow.construct(workflow_spec)
assert wf["add_watermark"].inputs.in_video == LazyField(
name=WORKFLOW_LZIN, field="input_video", type=video.Mp4, type_checked=True
)
assert wf["add_watermark"].inputs.watermark == LazyField(
name=WORKFLOW_LZIN, field="watermark", type=image.Png, type_checked=True
)


def test_direct_access_of_workflow_object():

@python.define(inputs={"x": float}, outputs={"z": float})
Expand Down
8 changes: 4 additions & 4 deletions pydra/engine/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -319,11 +319,11 @@ def construct(
for lzy_inpt in lazy_inputs:
setattr(
lazy_spec,
lzy_inpt,
lzy_inpt.name,
LazyField(
WORKFLOW_LZIN,
lzy_inpt.name,
lzy_inpt.type,
name=WORKFLOW_LZIN,
field=lzy_inpt.name,
type=lzy_inpt.type,
),
)

Expand Down

0 comments on commit a77be6a

Please sign in to comment.