Merge 75e3659 into 37d244e

Author: tobiasraabe

src/_pytask/collect_utils.py

@@ -308,9 +308,7 @@ 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)
-        ]
+        annot = [i for i in meta if not isinstance(i, ProductType)]
         if len(annot) >= 2:  # noqa: PLR2004
             raise ValueError(
                 f"Parameter {name!r} has multiple node annotations although only one "
@@ -360,6 +358,8 @@ def parse_products_from_task_function(
     has_produces_decorator = False
     has_produces_argument = False
     has_annotation = False
+    has_return = False
+    has_task_decorator = False
     out = {}
 
     # Parse products from decorators.
@@ -415,8 +415,45 @@ def parse_products_from_task_function(
                 )
                 out = {parameter_name: collected_products}
 
+    if "return" in parameters_with_node_annot:
+        has_return = True
+        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}
+
+    task_produces = obj.pytask_meta.produces if hasattr(obj, "pytask_meta") else None
+    if task_produces:
+        has_task_decorator = True
+        collected_products = tree_map_with_path(
+            lambda p, x: _collect_product(
+                session,
+                path,
+                name,
+                NodeInfo("return", p, x),
+                is_string_allowed=False,
+            ),
+            task_produces,
+        )
+        out = {"return": collected_products}
+
     if (
-        sum((has_produces_decorator, has_produces_argument, has_annotation))
+        sum(
+            (
+                has_produces_decorator,
+                has_produces_argument,
+                has_annotation,
+                has_return,
+                has_task_decorator,
+            )
+        )
         >= 2  # noqa: PLR2004
     ):
         raise NodeNotCollectedError(_ERROR_MULTIPLE_PRODUCT_DEFINITIONS)
@@ -571,5 +608,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.from_annot(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/models.py

@@ -5,6 +5,7 @@
 from typing import NamedTuple
 from typing import TYPE_CHECKING
 
+from _pytask.tree_util import PyTree
 from attrs import define
 from attrs import field
 
@@ -24,6 +25,8 @@ class CollectionMetadata:
     """Contains the markers of the function."""
     name: str | None = None
     """The name of the task function."""
+    produces: PyTree[Any] = None
+    """Definition of products to handle returns."""
 
 
 class NodeInfo(NamedTuple):

src/_pytask/node_protocols.py

@@ -15,7 +15,13 @@ class MetaNode(Protocol):
     """The name of node that must be unique."""
 
     @abstractmethod
-    def state(self) -> Any:
+    def state(self) -> str | None:
+        """Return the state of the node.
+
+        The state can be something like a hash or a last modified timestamp. If the node
+        does not exist, you can also return ``None``.
+
+        """
         ...
 
 
@@ -25,9 +31,36 @@ class Node(MetaNode, Protocol):
 
     value: Any
 
+    def load(self) -> Any:
+        """Return the value of the node that will be injected into the task."""
+        ...
+
+    def save(self, value: Any) -> Any:
+        """Save the value that was returned from a task."""
+        ...
+
+    def from_annot(self, value: Any) -> Any:
+        """Complete the node by setting the value from an default argument.
+
+        Use it, if you want to add information on how a node handles an argument while
+        keeping the type of the value unrelated to pytask.
+
+        .. codeblock: python
+
+            def task_example(value: Annotated[Any, PythonNode(hash=True)], produces):
+                ...
+
+
+        """
+        ...
+
 
 @runtime_checkable
 class PPathNode(Node, Protocol):
-    """Nodes with paths."""
+    """Nodes with paths.
+
+    Nodes with paths receive special handling when it comes to printing their names.
+
+    """
 
     path: Path

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
 
@@ -28,6 +31,7 @@ class ProductType:
 
 
 Product = ProductType()
+"""ProductType: A singleton to mark products in annotations."""
 
 
 @define(kw_only=True)
@@ -64,6 +68,7 @@ def __attrs_post_init__(self: Task) -> None:
             self.short_name = self.name
 
     def state(self, hash: bool = False) -> str | None:  # noqa: A002
+        """Return the state of the node."""
         if hash and self.path.exists():
             return hashlib.sha256(self.path.read_bytes()).hexdigest()
         if not hash and self.path.exists():
@@ -72,7 +77,22 @@ def state(self, hash: bool = False) -> str | None:  # noqa: A002
 
     def execute(self, **kwargs: Any) -> None:
         """Execute the task."""
-        self.function(**kwargs)
+        out = self.function(**kwargs)
+
+        if "return" in self.produces:
+            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 "
+                    f"of the return annotation.\n\nFunction return: {structure_out}\n\n"
+                    f"Return annotation: {structure_return}"
+                )
+
+            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."""
@@ -86,24 +106,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 from_annot(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
@@ -127,18 +143,45 @@ def state(self) -> str | None:
             return str(self.path.stat().st_mtime)
         return None
 
+    def load(self) -> Path:
+        """Load the value."""
+        return self.value
+
+    def save(self, value: bytes | str) -> None:
+        """Save strings or bytes to file."""
+        if isinstance(value, str):
+            self.path.write_text(value)
+        elif isinstance(value, bytes):
+            self.path.write_bytes(value)
+        else:
+            raise TypeError(
+                f"'PathNode' can only save 'str' and 'bytes', not {type(value)}"
+            )
+
 
 @define(kw_only=True)
 class PythonNode(Node):
     """The class for a node which is a Python object."""
 
     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:
+        """Load the value."""
+        return self.value
+
+    def save(self, value: Any) -> NoReturn:
+        """Save the value."""
+        raise NotImplementedError
+
+    def from_annot(self, value: Any) -> None:
+        """Set the value from a function annotation."""
+        self.value = value
+
     def state(self) -> str | None:
         """Calculate state of the node.
 

src/_pytask/task_utils.py

@@ -11,6 +11,7 @@
 from _pytask.mark import Mark
 from _pytask.models import CollectionMetadata
 from _pytask.shared import find_duplicates
+from _pytask.tree_util import PyTree
 
 
 __all__ = ["parse_keyword_arguments_from_signature_defaults"]
@@ -31,6 +32,7 @@ def task(
     *,
     id: str | None = None,  # noqa: A002
     kwargs: dict[Any, Any] | None = None,
+    produces: PyTree[Any] = None,
 ) -> Callable[..., None]:
     """Parse inputs of the ``@pytask.mark.task`` decorator.
 
@@ -43,13 +45,15 @@ def task(
 
     Parameters
     ----------
-    name : str | None
+    name
         The name of the task.
-    id : str | None
+    id
         An id for the task if it is part of a parametrization.
-    kwargs : dict[Any, Any] | None
+    kwargs
         A dictionary containing keyword arguments which are passed to the task when it
         is executed.
+    produces
+        Definition of products to handle returns.
 
     """
 
@@ -77,12 +81,14 @@ def wrapper(func: Callable[..., Any]) -> None:
             unwrapped.pytask_meta.kwargs = parsed_kwargs
             unwrapped.pytask_meta.markers.append(Mark("task", (), {}))
             unwrapped.pytask_meta.id_ = id
+            unwrapped.pytask_meta.produces = produces
         else:
             unwrapped.pytask_meta = CollectionMetadata(
                 name=parsed_name,
                 kwargs=parsed_kwargs,
                 markers=[Mark("task", (), {})],
                 id_=id,
+                produces=produces,
             )
 
         COLLECTED_TASKS[path].append(unwrapped)

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()