Allow to inject task returns via @pytask.mark.task(produces=...).

Author: tobiasraabe

src/_pytask/collect_utils.py

@@ -359,6 +359,7 @@ def parse_products_from_task_function(
     has_produces_argument = False
     has_annotation = False
     has_return = False
+    has_task_decorator = False
     out = {}
 
     # Parse products from decorators.
@@ -428,8 +429,31 @@ def parse_products_from_task_function(
         )
         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, has_return))
+        sum(
+            (
+                has_produces_decorator,
+                has_produces_argument,
+                has_annotation,
+                has_return,
+                has_task_decorator,
+            )
+        )
         >= 2  # noqa: PLR2004
     ):
         raise NodeNotCollectedError(_ERROR_MULTIPLE_PRODUCT_DEFINITIONS)

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

tests/test_execute.py

@@ -585,6 +585,27 @@ def task_example() -> Annotated[str, node]:
     assert tmp_path.joinpath("file.txt").read_text() == "Hello, World!"
 
 
+@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!"
+
+
 @pytest.mark.end_to_end()
 def test_return_with_custom_type_annotation_as_return(runner, tmp_path):
     source = """
@@ -626,24 +647,3 @@ 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_task.py

@@ -511,3 +511,47 @@ def task_example(
     assert result.exit_code == ExitCode.OK
     assert tmp_path.joinpath("out.txt").read_text() == "Hello world!"
     assert not tmp_path.joinpath("not_used_out.txt").exists()
+
+
+@pytest.mark.end_to_end()
+def test_return_with_task_decorator(runner, tmp_path):
+    source = """
+    from pathlib import Path
+    from typing import Any
+    from typing_extensions import Annotated
+    from pytask import PathNode
+    import pytask
+
+    node = PathNode.from_path(Path(__file__).parent.joinpath("file.txt"))
+
+    @pytask.mark.task(produces=node)
+    def task_example():
+        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("file.txt").read_text() == "Hello, World!"
+
+
+@pytest.mark.end_to_end()
+def test_return_with_tuple_and_task_decorator(runner, tmp_path):
+    source = """
+    from pathlib import Path
+    from typing import Any
+    from typing_extensions import Annotated
+    from pytask import PathNode
+    import pytask
+
+    node1 = PathNode.from_path(Path(__file__).parent.joinpath("file1.txt"))
+    node2 = PathNode.from_path(Path(__file__).parent.joinpath("file2.txt"))
+
+    @pytask.mark.task(produces=(node1, node2))
+    def task_example():
+        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!"