Add is_product to PNode.load(). (#472)

Author: tobiasraabe

.pre-commit-config.yaml

@@ -51,6 +51,7 @@ repos:
     rev: v0.1.1
     hooks:
       - id: ruff
+        args: [--unsafe-fixes]
 -   repo: https://github.com/dosisod/refurb
     rev: v1.22.1
     hooks:

docs/source/changes.md

@@ -19,6 +19,7 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
 - {pull}`463` raise error when a task function is not defined inside the loop body.
 - {pull}`464` improves pinned dependencies.
 - {pull}`465` adds test to ensure internal tracebacks are removed by reports.
+- {pull}`472` adds `is_product` to {meth}`PNode.load`.
 
 ## 0.4.1 - 2023-10-11
 

docs/source/how_to_guides/writing_custom_nodes.md

@@ -20,9 +20,6 @@ to inputs and outputs and call {func}`pandas.read_pickle` and
 To remove IO operations from the task and delegate them to pytask, we will write a
 `PickleNode` that automatically loads and stores Python objects.
 
-We will also use the feature explained in {doc}`using_task_returns` to define products
-of the task function via the function's return value.
-
 And we pass the value to `df` via {obj}`Annotated` to preserve the type hint.
 
 The result will be the following task.
@@ -37,12 +34,28 @@ The result will be the following task.
 
 :::
 
+:::{tab-item} Python 3.10+ & Return
+:sync: python310plus
+
+```{literalinclude} ../../../docs_src/how_to_guides/writing_custom_nodes_example_2_py310_return.py
+```
+
+:::
+
 :::{tab-item} Python 3.8+
 :sync: python38plus
 
 ```{literalinclude} ../../../docs_src/how_to_guides/writing_custom_nodes_example_2_py38.py
 ```
 
+:::
+
+:::{tab-item} Python 3.8+ & Return
+:sync: python38plus
+
+```{literalinclude} ../../../docs_src/how_to_guides/writing_custom_nodes_example_2_py38_return.py
+```
+
 :::
 ::::
 
@@ -97,7 +110,12 @@ Here are some explanations.
   the value changes, pytask knows it needs to regenerate the workflow. We can use
   the timestamp of when the node was last modified.
 - pytask calls {meth}`PickleNode.load` when it collects the values of function arguments
-  to run the function. In our example, we read the file and unpickle the data.
+  to run the function. The argument `is_product` signals that the node is loaded as a
+  product with a {class}`~pytask.Product` annotation or via `produces`.
+
+  When the node is loaded as a dependency, we want to inject the value of the pickle
+  file. In the other case, the node returns itself so users can call
+  {meth}`PickleNode.save` themselves.
 - {meth}`PickleNode.save` is called when a task function returns and allows to save the
   return values.
 

docs_src/how_to_guides/writing_custom_nodes_example_2_py310.py

@@ -2,6 +2,7 @@
 from typing import Annotated
 
 import pandas as pd
+from pytask import Product
 
 
 class PickleNode:
@@ -13,6 +14,7 @@ class PickleNode:
 
 
 def task_example(
-    df: Annotated[pd.DataFrame, in_node]
-) -> Annotated[pd.DataFrame, out_node]:
-    return df.apply(...)
+    df: Annotated[pd.DataFrame, in_node], out: Annotated[PickleNode, out_node, Product]
+) -> None:
+    transformed = df.apply(...)
+    out.save(transformed)

docs_src/how_to_guides/writing_custom_nodes_example_2_py310_return.py

@@ -0,0 +1,18 @@
+from pathlib import Path
+from typing import Annotated
+
+import pandas as pd
+
+
+class PickleNode:
+    ...
+
+
+in_node = PickleNode.from_path(Path(__file__).parent / "in.pkl")
+out_node = PickleNode.from_path(Path(__file__).parent / "out.pkl")
+
+
+def task_example(
+    df: Annotated[pd.DataFrame, in_node]
+) -> Annotated[pd.DataFrame, out_node]:
+    return df.apply(...)

docs_src/how_to_guides/writing_custom_nodes_example_2_py38.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import pandas as pd
+from pytask import Product
 from typing_extensions import Annotated
 
 
@@ -13,6 +14,7 @@ class PickleNode:
 
 
 def task_example(
-    df: Annotated[pd.DataFrame, in_node]
-) -> Annotated[pd.DataFrame, out_node]:
-    return df.apply(...)
+    df: Annotated[pd.DataFrame, in_node], out: Annotated[PickleNode, out_node, Product]
+) -> None:
+    transformed = df.apply(...)
+    out.save(transformed)

docs_src/how_to_guides/writing_custom_nodes_example_2_py38_return.py

@@ -0,0 +1,18 @@
+from pathlib import Path
+
+import pandas as pd
+from typing_extensions import Annotated
+
+
+class PickleNode:
+    ...
+
+
+in_node = PickleNode.from_path(Path(__file__).parent / "in.pkl")
+out_node = PickleNode.from_path(Path(__file__).parent / "out.pkl")
+
+
+def task_example(
+    df: Annotated[pd.DataFrame, in_node]
+) -> Annotated[pd.DataFrame, out_node]:
+    return df.apply(...)

docs_src/how_to_guides/writing_custom_nodes_example_3_py310.py

@@ -33,8 +33,10 @@ def state(self) -> str | None:
             return str(self.path.stat().st_mtime)
         return None
 
-    def load(self) -> Path:
+    def load(self, is_product: bool) -> Path:
         """Load the value from the path."""
+        if is_product:
+            return self
         return pickle.loads(self.path.read_bytes())
 
     def save(self, value: Any) -> None:

docs_src/how_to_guides/writing_custom_nodes_example_3_py38.py

@@ -34,8 +34,10 @@ def state(self) -> Optional[str]:
             return str(self.path.stat().st_mtime)
         return None
 
-    def load(self) -> Path:
+    def load(self, is_product: bool) -> Path:
         """Load the value from the path."""
+        if is_product:
+            return self
         return pickle.loads(self.path.read_bytes())
 
     def save(self, value: Any) -> None:

src/_pytask/execute.py

@@ -144,9 +144,9 @@ def pytask_execute_task_setup(session: Session, task: PTask) -> None:
         raise WouldBeExecuted
 
 
-def _safe_load(node: PNode, task: PTask) -> Any:
+def _safe_load(node: PNode, task: PTask, is_product: bool) -> Any:
     try:
-        return node.load()
+        return node.load(is_product=is_product)
     except Exception as e:  # noqa: BLE001
         task_name = getattr(task, "display_name", task.name)
         msg = f"Exception while loading node {node.name!r} of task {task_name!r}"
@@ -163,11 +163,11 @@ def pytask_execute_task(session: Session, task: PTask) -> bool:
 
     kwargs = {}
     for name, value in task.depends_on.items():
-        kwargs[name] = tree_map(lambda x: _safe_load(x, task), value)
+        kwargs[name] = tree_map(lambda x: _safe_load(x, task, False), value)
 
     for name, value in task.produces.items():
         if name in parameters:
-            kwargs[name] = tree_map(lambda x: _safe_load(x, task), value)
+            kwargs[name] = tree_map(lambda x: _safe_load(x, task, True), value)
 
     out = task.execute(**kwargs)
 

src/_pytask/node_protocols.py

@@ -39,8 +39,15 @@ def state(self) -> str | None:
 class PNode(MetaNode, Protocol):
     """Protocol for nodes."""
 
-    def load(self) -> Any:
-        """Return the value of the node that will be injected into the task."""
+    def load(self, is_product: bool) -> Any:
+        """Return the value of the node that will be injected into the task.
+
+        Parameters
+        ----------
+        is_product
+            Indicates whether the node is loaded as a dependency or as a product.
+
+        """
         ...
 
     def save(self, value: Any) -> Any:

src/_pytask/nodes.py

@@ -173,7 +173,7 @@ def state(self) -> str | None:
             return str(self.path.stat().st_mtime)
         return None
 
-    def load(self) -> Path:
+    def load(self, is_product: bool = False) -> Path:  # noqa: ARG002
         """Load the value."""
         return self.path
 
@@ -207,8 +207,10 @@ class PythonNode(PNode):
     value: Any | NoDefault = no_default
     hash: bool | Callable[[Any], bool] = False  # noqa: A003
 
-    def load(self) -> Any:
+    def load(self, is_product: bool = False) -> Any:
         """Load the value."""
+        if is_product:
+            return self
         if isinstance(self.value, PythonNode):
             return self.value.load()
         return self.value

tests/test_execute.py

@@ -606,7 +606,7 @@ def task_example() -> Annotated[str, (node1, node2)]:
 
 
 @pytest.mark.end_to_end()
-def test_return_with_custom_type_annotation_as_return(runner, tmp_path):
+def test_return_with_custom_node_and_return_annotation(runner, tmp_path):
     source = """
     from __future__ import annotations
 
@@ -626,7 +626,9 @@ def state(self) -> str | None:
                 return str(self.path.stat().st_mtime)
             return None
 
-        def load(self) -> Any:
+        def load(self, is_product) -> Any:
+            if is_product:
+                return self
             return pickle.loads(self.path.read_bytes())
 
         def save(self, value: Any) -> None:
@@ -645,6 +647,49 @@ def task_example() -> Annotated[int, node]:
     assert data == 1
 
 
+@pytest.mark.end_to_end()
+def test_return_with_custom_node_with_product_annotation(runner, tmp_path):
+    source = """
+    from __future__ import annotations
+
+    from pathlib import Path
+    import pickle
+    from typing import Any
+    from typing_extensions import Annotated
+    import attrs
+    from pytask import Product
+
+    @attrs.define
+    class PickleNode:
+        name: str
+        path: Path
+
+        def state(self) -> str | None:
+            if self.path.exists():
+                return str(self.path.stat().st_mtime)
+            return None
+
+        def load(self, is_product) -> Any:
+            if is_product:
+                return self
+            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(node: Annotated[PickleNode, node, Product]) -> None:
+        node.save(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
+
+
 @pytest.mark.end_to_end()
 def test_error_when_return_pytree_mismatch(runner, tmp_path):
     source = """
@@ -835,3 +880,23 @@ def task_example(
 
     # Test that traceback is hidden.
     assert "_pytask/execute.py" not in result.output
+
+
+def test_python_node_as_product_with_product_annotation(runner, tmp_path):
+    source = """
+    from typing_extensions import Annotated
+    from pytask import Product, PythonNode
+    from pathlib import Path
+
+    node = PythonNode()
+
+    def task_create_string(node: Annotated[PythonNode, node, Product]) -> None:
+        node.save("Hello, World!")
+
+    def task_write_file(text: Annotated[str, node]) -> Annotated[str, Path("file.txt")]:
+        return text
+    """
+    tmp_path.joinpath("task_example.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!"

tests/test_node_protocols.py

@@ -24,7 +24,7 @@ class CustomNode:
         def state(self):
             return self.value
 
-        def load(self):
+        def load(self, is_product):
             return self.value
 
         def save(self, value):
@@ -61,7 +61,7 @@ class PickleFile:
         def state(self):
             return str(self.path.stat().st_mtime)
 
-        def load(self):
+        def load(self, is_product):
             with self.path.open("rb") as f:
                 out = pickle.load(f)
             return out

tests/test_nodes.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import pytest
-from _pytask.nodes import PythonNode
+from pytask import PathNode
+from pytask import PNode
+from pytask import PPathNode
+from pytask import PythonNode
 
 
 @pytest.mark.unit()
@@ -19,3 +22,16 @@ def test_hash_of_python_node(value, hash_, expected):
     node = PythonNode(name="test", value=value, hash=hash_)
     state = node.state()
     assert state == expected
+
+
+@pytest.mark.parametrize(
+    ("node", "protocol", "expected"),
+    [
+        (PathNode, PNode, True),
+        (PathNode, PPathNode, True),
+        (PythonNode, PNode, True),
+        (PythonNode, PPathNode, False),
+    ],
+)
+def test_comply_with_protocol(node, protocol, expected):
+    assert isinstance(node, protocol) is expected