Allow pytrees in return annotations and function returns.

Author: tobiasraabe

src/_pytask/collect_utils.py

@@ -309,7 +309,10 @@ def _find_args_with_node_annotation(func: Callable[..., Any]) -> dict[str, Node]
     args_with_node_annotation = {}
     for name, meta in metas.items():
         annot = [
-            i for i in meta if not isinstance(i, ProductType) and isinstance(i, Node)
+            i
+            for i in meta
+            if not isinstance(i, ProductType)
+            and all(isinstance(x, Node) for x in tree_leaves(i))
         ]
         if len(annot) >= 2:  # noqa: PLR2004
             raise ValueError(
@@ -418,7 +421,17 @@ def parse_products_from_task_function(
 
     if "return" in parameters_with_node_annot:
         has_return = True
-        out = {"return": parameters_with_node_annot["return"]}
+        collected_products = tree_map_with_path(
+            lambda p, x: _collect_product(
+                session,
+                path,
+                name,
+                NodeInfo("return", p, x),
+                is_string_allowed=False,
+            ),
+            parameters_with_node_annot["return"],
+        )
+        out = {"return": collected_products}
 
     if (
         sum((has_produces_decorator, has_produces_argument, has_annotation, has_return))
@@ -576,5 +589,5 @@ def _evolve_instance(x: Any, instance_from_annot: Node | None) -> Any:
     if not instance_from_annot:
         return x
 
-    instance_from_annot.value = x
+    instance_from_annot.set_value(x)
     return instance_from_annot

src/_pytask/execute.py

@@ -147,11 +147,11 @@ def pytask_execute_task(session: Session, task: Task) -> bool:
 
     kwargs = {}
     for name, value in task.depends_on.items():
-        kwargs[name] = tree_map(lambda x: x.value, value)
+        kwargs[name] = tree_map(lambda x: x.load(), value)
 
     for name, value in task.produces.items():
         if name in parameters:
-            kwargs[name] = tree_map(lambda x: x.value, value)
+            kwargs[name] = tree_map(lambda x: x.load(), value)
 
     task.execute(**kwargs)
     return True

src/_pytask/node_protocols.py

@@ -25,6 +25,15 @@ class Node(MetaNode, Protocol):
 
     value: Any
 
+    def load(self) -> Any:
+        ...
+
+    def save(self, value: Any) -> Any:
+        ...
+
+    def set_value(self, value: Any) -> Any:
+        ...
+
 
 @runtime_checkable
 class PPathNode(Node, Protocol):

src/_pytask/nodes.py

@@ -6,11 +6,14 @@
 from pathlib import Path
 from typing import Any
 from typing import Callable
+from typing import NoReturn
 from typing import TYPE_CHECKING
 
 from _pytask.node_protocols import MetaNode
 from _pytask.node_protocols import Node
 from _pytask.tree_util import PyTree
+from _pytask.tree_util import tree_leaves
+from _pytask.tree_util import tree_structure
 from attrs import define
 from attrs import field
 
@@ -76,7 +79,18 @@ def execute(self, **kwargs: Any) -> None:
         out = self.function(**kwargs)
 
         if "return" in self.produces:
-            self.produces["return"].save(out)
+            structure_out = tree_structure(out)
+            structure_return = tree_structure(self.produces["return"])
+            if not structure_out == structure_return:
+                raise ValueError(
+                    "The structure of the function return does not match the structure "
+                    "of the return annotation."
+                )
+
+            for out_, return_ in zip(
+                tree_leaves(out), tree_leaves(self.produces["return"])
+            ):
+                return_.save(out_)
 
     def add_report_section(self, when: str, key: str, content: str) -> None:
         """Add sections which will be displayed in report like stdout or stderr."""
@@ -90,25 +104,20 @@ class PathNode(Node):
 
     name: str = ""
     """Name of the node which makes it identifiable in the DAG."""
-    _value: Path | None = None
+    value: Path | None = None
     """Value passed to the decorator which can be requested inside the function."""
 
     @property
     def path(self) -> Path:
         return self.value
 
-    @property
-    def value(self) -> Path:
-        return self._value
-
-    @value.setter
-    def value(self, value: Path) -> None:
+    def set_value(self, value: Path) -> None:
         """Set path and if other attributes are not set, set sensible defaults."""
         if not isinstance(value, Path):
             raise TypeError("'value' must be a 'pathlib.Path'.")
         if not self.name:
             self.name = value.as_posix()
-        self._value = value
+        self.value = value
 
     @classmethod
     @functools.lru_cache
@@ -153,11 +162,20 @@ class PythonNode(Node):
 
     name: str = ""
     """Name of the node."""
-    value: Any | None = None
+    value: Any = None
     """Value of the node."""
     hash: bool = False  # noqa: A003
     """Whether the value should be hashed to determine the state."""
 
+    def load(self) -> Any:
+        return self.value
+
+    def save(self, value: Any) -> NoReturn:
+        raise NotImplementedError
+
+    def set_value(self, value: Any) -> None:
+        self.value = value
+
     def state(self) -> str | None:
         """Calculate state of the node.
 

src/_pytask/tree_util.py

@@ -9,12 +9,14 @@
 from optree import tree_leaves as _optree_tree_leaves
 from optree import tree_map as _optree_tree_map
 from optree import tree_map_with_path as _optree_tree_map_with_path
+from optree import tree_structure as _optree_tree_structure
 
 
 __all__ = [
     "tree_leaves",
     "tree_map",
     "tree_map_with_path",
+    "tree_structure",
     "PyTree",
     "TREE_UTIL_LIB_DIRECTORY",
 ]
@@ -29,3 +31,6 @@
 tree_map_with_path = functools.partial(
     _optree_tree_map_with_path, none_is_leaf=True, namespace="pytask"
 )
+tree_structure = functools.partial(
+    _optree_tree_structure, none_is_leaf=True, namespace="pytask"
+)

tests/test_collect.py

@@ -172,7 +172,7 @@ def test_pytask_collect_node(session, path, node_info, expected):
     if result is None:
         assert result is expected
     else:
-        assert str(result.value) == str(expected)
+        assert str(result.load()) == str(expected)
 
 
 @pytest.mark.unit()

tests/test_collect_command.py

@@ -529,6 +529,9 @@ class CustomNode:
         def state(self):
             return self.value
 
+        def load(self): ...
+        def save(self, value): ...
+        def set_value(self, value): ...
 
     def task_example(
         data = CustomNode("custom", "text"),
@@ -555,21 +558,27 @@ def test_node_protocol_for_custom_nodes_with_paths(runner, tmp_path):
     class PickleFile:
         name: str
         path: Path
+        value: Path
 
-        @property
-        def value(self):
+        def state(self):
+            return str(self.path.stat().st_mtime)
+
+        def load(self):
             with self.path.open("rb") as f:
                 out = pickle.load(f)
             return out
 
-        def state(self):
-            return str(self.path.stat().st_mtime)
+        def save(self, value):
+            with self.path.open("wb") as f:
+                pickle.dump(value, f)
+
+        def set_value(self, value): ...
 
 
     _PATH = Path(__file__).parent.joinpath("in.pkl")
 
     def task_example(
-        data = PickleFile(_PATH.as_posix(), _PATH),
+        data = PickleFile(_PATH.as_posix(), _PATH, _PATH),
         out: Annotated[Path, Product] = Path("out.txt"),
     ) -> None:
         out.write_text(data)

tests/test_execute.py

@@ -532,6 +532,7 @@ def test_error_with_multiple_different_dep_annotations(runner, tmp_path):
     from pathlib import Path
     from typing_extensions import Annotated
     from pytask import Product, PythonNode, PathNode
+    from typing import Any
 
     def task_example(
         dependency: Annotated[Any, PythonNode(), PathNode()] = "hello",
@@ -591,6 +592,8 @@ def load(self) -> Any:
         def save(self, value: Any) -> None:
             self.path.write_bytes(pickle.dumps(value))
 
+        def set_value(self, value: Any) -> None: ...
+
     node = PickleNode("pickled_data", Path(__file__).parent.joinpath("data.pkl"))
 
     def task_example() -> Annotated[int, node]:
@@ -602,3 +605,24 @@ def task_example() -> Annotated[int, node]:
 
     data = pickle.loads(tmp_path.joinpath("data.pkl").read_bytes())  # noqa: S301
     assert data == 1
+
+
+@pytest.mark.end_to_end()
+def test_return_with_tuple_pathnode_annotation_as_return(runner, tmp_path):
+    source = """
+    from pathlib import Path
+    from typing import Any
+    from typing_extensions import Annotated
+    from pytask import PathNode
+
+    node1 = PathNode.from_path(Path(__file__).parent.joinpath("file1.txt"))
+    node2 = PathNode.from_path(Path(__file__).parent.joinpath("file2.txt"))
+
+    def task_example() -> Annotated[str, (node1, node2)]:
+        return "Hello,", "World!"
+    """
+    tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source))
+    result = runner.invoke(cli, [tmp_path.as_posix()])
+    assert result.exit_code == ExitCode.OK
+    assert tmp_path.joinpath("file1.txt").read_text() == "Hello,"
+    assert tmp_path.joinpath("file2.txt").read_text() == "World!"

tests/test_node_protocols.py

@@ -22,6 +22,14 @@ class CustomNode:
         def state(self):
             return self.value
 
+        def load(self):
+            return self.value
+
+        def save(self, value):
+            self.value = value
+
+        def set_value(self, value): ...
+
 
     def task_example(
         data = CustomNode("custom", "text"),
@@ -48,21 +56,27 @@ def test_node_protocol_for_custom_nodes_with_paths(runner, tmp_path):
     class PickleFile:
         name: str
         path: Path
+        value: Path
+
+        def state(self):
+            return str(self.path.stat().st_mtime)
 
-        @property
-        def value(self):
+        def load(self):
             with self.path.open("rb") as f:
                 out = pickle.load(f)
             return out
 
-        def state(self):
-            return str(self.path.stat().st_mtime)
+        def save(self, value):
+            with self.path.open("wb") as f:
+                pickle.dump(value, f)
+
+        def set_value(self, value): ...
 
 
     _PATH = Path(__file__).parent.joinpath("in.pkl")
 
     def task_example(
-        data = PickleFile(_PATH.as_posix(), _PATH),
+        data = PickleFile(_PATH.as_posix(), _PATH, _PATH),
         out: Annotated[Path, Product] = Path("out.txt"),
     ) -> None:
         out.write_text(data)

tests/test_tree_util.py

@@ -43,7 +43,7 @@ def task_example():
 
     assert session.exit_code == exit_code
 
-    products = tree_map(lambda x: x.value, getattr(session.tasks[0], decorator_name))
+    products = tree_map(lambda x: x.load(), getattr(session.tasks[0], decorator_name))
     expected = {
         0: tmp_path / "out.txt",
         1: {0: tmp_path / "tuple_out.txt"},