fixed up retrospective setting of node inputs via the inputs attribute of the TaskOutputs object

Author: tclose

pydra/compose/base/task.py

@@ -39,11 +39,16 @@ class Outputs:
     RESERVED_FIELD_NAMES = ("inputs",)
 
     _cache_dir: Path = attrs.field(default=None, init=False, repr=False)
+    _node = attrs.field(default=None, init=False, repr=False)
 
     @property
     def inputs(self):
         """The inputs object associated with a lazy-outputs object"""
-        return self._get_node().inputs
+        if self._node is None:
+            raise AttributeError(
+                f"{self} outputs object is not a lazy output of a workflow node"
+            )
+        return self._node.inputs
 
     @classmethod
     def _from_task(cls, job: "Job[TaskType]") -> Self:
@@ -81,14 +86,6 @@ def _results(self) -> "Result[Self]":
         with open(results_path, "rb") as f:
             return cp.load(f)
 
-    def _get_node(self):
-        try:
-            return self._node
-        except AttributeError:
-            raise AttributeError(
-                f"{self} outputs object is not a lazy output of a workflow node"
-            ) from None
-
     def __iter__(self) -> ty.Generator[str, None, None]:
         """The names of the fields in the output object"""
         return iter(sorted(f.name for f in attrs_fields(self)))

pydra/compose/tests/test_workflow_fields.py

@@ -1,4 +1,5 @@
 from operator import attrgetter
+from pathlib import Path
 from copy import copy
 from unittest.mock import Mock
 import pytest
@@ -15,17 +16,17 @@
 
 
 @python.define
-def Add(a, b):
+def Add(a: int | float, b: int | float) -> int | float:
     return a + b
 
 
 @python.define
-def Mul(a, b):
+def Mul(a: int | float, b: int | float) -> int | float:
     return a * b
 
 
 @python.define(outputs=["divided"])
-def Divide(x, y):
+def Divide(x: int | float, y: int | float) -> float:
     return x / y
 
 
@@ -68,7 +69,9 @@ def MyTestWorkflow(a, b):
     wf = Workflow.construct(workflow_spec)
     assert wf.inputs.a == 1
     assert wf.inputs.b == 2.0
-    assert wf.outputs.out == LazyOutField(node=wf["Mul"], field="out", type=ty.Any)
+    assert wf.outputs.out == LazyOutField(
+        node=wf["Mul"], field="out", type=int | float, type_checked=True
+    )
 
     # Nodes are named after the specs by default
     assert list(wf.node_names) == ["Add", "Mul"]
@@ -185,7 +188,9 @@ class Outputs(workflow.Outputs):
     wf = Workflow.construct(workflow_spec)
     assert wf.inputs.a == 1
     assert wf.inputs.b == 2.0
-    assert wf.outputs.out == LazyOutField(node=wf["Mul"], field="out", type=ty.Any)
+    assert wf.outputs.out == LazyOutField(
+        node=wf["Mul"], field="out", type=int | float, type_checked=True
+    )
 
     # Nodes are named after the specs by default
     assert list(wf.node_names) == ["Add", "Mul"]
@@ -323,7 +328,7 @@ def MyTestWorkflow(a: int, b: float) -> tuple[float, float]:
         node=wf["Mul"], field="out", type=float, type_checked=True
     )
     assert wf.outputs.out2 == LazyOutField(
-        node=wf["division"], field="divided", type=ty.Any
+        node=wf["division"], field="divided", type=float, type_checked=True
     )
     assert list(wf.node_names) == ["addition", "Mul", "division"]
 
@@ -362,8 +367,12 @@ def MyTestWorkflow(a: int, b: float):
     wf = Workflow.construct(workflow_spec)
     assert wf.inputs.a == 1
     assert wf.inputs.b == 2.0
-    assert wf.outputs.out1 == LazyOutField(node=wf["Mul"], field="out", type=ty.Any)
-    assert wf.outputs.out2 == LazyOutField(node=wf["Add"], field="out", type=ty.Any)
+    assert wf.outputs.out1 == LazyOutField(
+        node=wf["Mul"], field="out", type=int | float, type_checked=True
+    )
+    assert wf.outputs.out2 == LazyOutField(
+        node=wf["Add"], field="out", type=int | float, type_checked=True
+    )
     assert list(wf.node_names) == ["Add", "Mul"]
 
 
@@ -500,3 +509,68 @@ def RecursiveNestedWorkflow(a: float, depth: int) -> float:
         type=float,
         type_checked=True,
     )
+
+
+def test_workflow_lzout_inputs1(tmp_path: Path):
+
+    @workflow.define
+    def InputAccessWorkflow(a, b, c):
+        add = workflow.add(Add(a=a, b=b))
+        add.inputs.a = c
+        mul = workflow.add(Mul(a=add.out, b=b))
+        return mul.out
+
+    input_access_workflow = InputAccessWorkflow(a=1, b=2.0, c=3.0)
+    outputs = input_access_workflow(cache_root=tmp_path)
+    assert outputs.out == 10.0
+
+
+def test_workflow_lzout_inputs2(tmp_path: Path):
+
+    @workflow.define
+    def InputAccessWorkflow(a, b, c):
+        add = workflow.add(Add(a=a, b=b))
+        add.inputs.a = c
+        mul = workflow.add(Mul(a=add.out, b=b))
+        return mul.out
+
+    input_access_workflow = InputAccessWorkflow(a=1, b=2.0, c=3.0)
+    outputs = input_access_workflow(cache_root=tmp_path)
+    assert outputs.out == 10.0
+
+
+def test_workflow_lzout_inputs2(tmp_path: Path):
+    """Set the inputs of the 'add' node after its outputs have been accessed
+    but the state has not been altered"""
+
+    @workflow.define
+    def InputAccessWorkflow(a, b, c):
+        add = workflow.add(Add(a=a, b=b))
+        mul = workflow.add(Mul(a=add.out, b=b))
+        add.inputs.a = c
+        return mul.out
+
+    input_access_workflow = InputAccessWorkflow(a=1, b=2.0, c=3.0)
+    outputs = input_access_workflow(cache_root=tmp_path)
+    assert outputs.out == 10.0
+
+
+def test_workflow_lzout_inputs_state_change_fail(tmp_path: Path):
+    """Set the inputs of the 'mul' node after its outputs have been accessed
+    with an upstream lazy field that has a different state than the original.
+    This changes the type of the input and is therefore not permitted"""
+
+    @workflow.define
+    def InputAccessWorkflow(a, b, c):
+        add1 = workflow.add(Add(a=a, b=b), name="add1")
+        add2 = workflow.add(Add(a=a).split(b=c), name="add2")
+        mul1 = workflow.add(Mul(a=add1.out, b=b), name="mul1")
+        mul2 = workflow.add(Mul(a=mul1.out, b=b), name="mul2")
+        mul1.inputs.a = add2.out
+        return mul2.out
+
+    input_access_workflow = InputAccessWorkflow(a=1, b=2.0, c=[3.0, 4.0])
+    with pytest.raises(
+        RuntimeError, match="have already been accessed and therefore cannot set"
+    ):
+        input_access_workflow.construct()

pydra/engine/node.py

@@ -3,7 +3,7 @@
 from enum import Enum
 import attrs
 from pydra.engine import lazy
-from pydra.utils.general import attrs_values
+from pydra.utils.general import attrs_values, task_dict
 from pydra.utils.typing import is_lazy
 from pydra.engine.state import State, add_name_splitter, add_name_combiner
 
@@ -79,7 +79,7 @@ def __setattr__(self, name: str, value: ty.Any) -> None:
                         f"cannot set {name!r} input to {value} because it changes the "
                         f"state"
                     )
-                    self._set_state()
+                    self._node._set_state()
 
     @property
     def inputs(self) -> Inputs:
@@ -144,6 +144,7 @@ def lzout(self) -> OutputType:
             # output of an upstream node with additional state variables.
             outpt._type_checked = False
         self._lzout = outputs
+        outputs._node = self
         return outputs
 
     @property
@@ -161,10 +162,8 @@ def combiner(self):
     def _check_if_outputs_have_been_used(self, msg):
         used = []
         if self._lzout:
-            for outpt_name, outpt_val in attrs.asdict(
-                self._lzout, recurse=False
-            ).items():
-                if outpt_val.type_checked:
+            for outpt_name, outpt_val in task_dict(self._lzout).items():
+                if outpt_val._type_checked:
                     used.append(outpt_name)
         if used:
             raise RuntimeError(