Add test showing how to parse custom products from path annotations while introducing new load and save methods.

Author: tobiasraabe

src/_pytask/collect_utils.py

@@ -360,6 +360,7 @@ def parse_products_from_task_function(
     has_produces_decorator = False
     has_produces_argument = False
     has_annotation = False
+    has_return = False
     out = {}
 
     # Parse products from decorators.
@@ -415,8 +416,12 @@ def parse_products_from_task_function(
                 )
                 out = {parameter_name: collected_products}
 
+    if "return" in parameters_with_node_annot:
+        has_return = True
+        out = {"return": parameters_with_node_annot["return"]}
+
     if (
-        sum((has_produces_decorator, has_produces_argument, has_annotation))
+        sum((has_produces_decorator, has_produces_argument, has_annotation, has_return))
         >= 2  # noqa: PLR2004
     ):
         raise NodeNotCollectedError(_ERROR_MULTIPLE_PRODUCT_DEFINITIONS)

src/_pytask/nodes.py

@@ -72,7 +72,10 @@ 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:
+            self.produces["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."""
@@ -127,6 +130,9 @@ def state(self) -> str | None:
             return str(self.path.stat().st_mtime)
         return None
 
+    def load(self) -> Path:
+        return self.value
+
 
 @define(kw_only=True)
 class PythonNode(Node):

tests/test_execute.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import pickle
 import re
 import subprocess
 import sys
@@ -544,3 +545,42 @@ def task_example(
     result = runner.invoke(cli, [tmp_path.as_posix()])
     assert result.exit_code == ExitCode.COLLECTION_FAILED
     assert "Parameter 'dependency'" in result.output
+
+
+@pytest.mark.end_to_end()
+def test_return_with_custom_type_annotation_as_return(runner, tmp_path):
+    source = """
+    from pathlib import Path
+    import pickle
+    from typing import Any
+    from typing_extensions import Annotated
+    import attrs
+
+    @attrs.define
+    class PickleNode:
+        name: str = ""
+        path: Path | None = None
+        value: None = None
+
+        def state(self) -> str | None:
+            if self.path.exists():
+                return str(self.path.stat().st_mtime)
+            return None
+
+        def load(self) -> Any:
+            return pickle.loads(self.path.read_bytes())
+
+        def save(self, value: Any) -> None:
+            self.path.write_bytes(pickle.dumps(value))
+
+    node = PickleNode("pickled_data", Path(__file__).parent.joinpath("data.pkl"))
+
+    def task_example() -> Annotated[int, node]:
+        return 1
+    """
+    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
+
+    data = pickle.loads(tmp_path.joinpath("data.pkl").read_bytes())  # noqa: S301
+    assert data == 1