ENH: Make attributes mandatory on PNode and PProvisionalNode (#787)

Author: tobiasraabe

CHANGELOG.md

@@ -7,6 +7,9 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
 
 ## Unreleased
 
+- {pull}`787` makes the `attributes` field mandatory on `PNode` and
+  `PProvisionalNode`, and preserves existing node attributes when loading entries from
+  the data catalog.
 - {pull}`744` Removed the direct dependency on attrs and migrated internal models to
   dataclasses.
 - {pull}`766` moves runtime profiling persistence from SQLite to a JSON snapshot plus

src/_pytask/collect.py

@@ -33,7 +33,6 @@
 from _pytask.node_protocols import PPathNode
 from _pytask.node_protocols import PProvisionalNode
 from _pytask.node_protocols import PTask
-from _pytask.node_protocols import warn_about_upcoming_attributes_field_on_nodes
 from _pytask.nodes import DirectoryNode
 from _pytask.nodes import PathNode
 from _pytask.nodes import PythonNode
@@ -419,9 +418,6 @@ def pytask_collect_node(  # noqa: C901, PLR0912
     """
     node = node_info.value
 
-    if isinstance(node, (PNode, PProvisionalNode)) and not hasattr(node, "attributes"):
-        warn_about_upcoming_attributes_field_on_nodes()
-
     if isinstance(node, DirectoryNode):
         if node.root_dir is None:
             node.root_dir = path

src/_pytask/data_catalog.py

@@ -21,7 +21,6 @@
 from _pytask.models import NodeInfo
 from _pytask.node_protocols import PNode
 from _pytask.node_protocols import PProvisionalNode
-from _pytask.node_protocols import warn_about_upcoming_attributes_field_on_nodes
 from _pytask.nodes import PickleNode
 from _pytask.pluginmanager import storage
 from _pytask.session import Session
@@ -102,10 +101,7 @@ def __post_init__(self) -> None:
         # Initialize the data catalog with persisted nodes from previous runs.
         for path in self.path.glob("*-node.pkl"):
             node = pickle.loads(path.read_bytes())  # noqa: S301
-            if not hasattr(node, "attributes"):
-                warn_about_upcoming_attributes_field_on_nodes()
-            else:
-                node.attributes = {DATA_CATALOG_NAME_FIELD: self.name}
+            node.attributes[DATA_CATALOG_NAME_FIELD] = self.name
             self._entries[node.name] = node
 
     def __getitem__(self, name: str) -> PNode | PProvisionalNode:
@@ -150,7 +146,4 @@ def add(self, name: str, node: PNode | PProvisionalNode | Any = None) -> None:
             self._entries[name] = collected_node
 
         node = self._entries[name]
-        if hasattr(node, "attributes"):
-            node.attributes[DATA_CATALOG_NAME_FIELD] = self.name  # ty: ignore[invalid-assignment]
-        else:
-            warn_about_upcoming_attributes_field_on_nodes()
+        node.attributes[DATA_CATALOG_NAME_FIELD] = self.name

src/_pytask/node_protocols.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import warnings
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Protocol
@@ -22,6 +21,7 @@ class PNode(Protocol):
     """Protocol for nodes."""
 
     name: str
+    attributes: dict[Any, Any]
 
     @property
     def signature(self) -> str:
@@ -117,6 +117,7 @@ class PProvisionalNode(Protocol):
     """
 
     name: str
+    attributes: dict[Any, Any]
 
     @property
     def signature(self) -> str:
@@ -139,14 +140,3 @@ def load(self, is_product: bool = False) -> Any:  # pragma: no cover
 
     def collect(self) -> list[Any]:
         """Collect the objects that are defined by the provisional nodes."""
-
-
-def warn_about_upcoming_attributes_field_on_nodes() -> None:
-    warnings.warn(
-        "PNode and PProvisionalNode will require an 'attributes' field starting "
-        "with pytask v0.6.0. It is a dictionary with any type of key and values "
-        "similar to PTask. See https://tinyurl.com/pytask-custom-nodes for more "
-        "information about adjusting your custom nodes.",
-        stacklevel=1,
-        category=FutureWarning,
-    )

tests/test_collect_command.py

@@ -468,15 +468,18 @@ def task_example(
 def test_node_protocol_for_custom_nodes(runner, tmp_path):
     source = """
     from typing import Annotated
+    from typing import Any
     from pytask import Product
     from dataclasses import dataclass
+    from dataclasses import field
     from pathlib import Path
 
     @dataclass
     class CustomNode:
         name: str
         value: str
         signature: str = "id"
+        attributes: dict[Any, Any] = field(default_factory=dict)
 
         def state(self):
             return self.value
@@ -673,9 +676,11 @@ def task_example(
 def test_collect_custom_node_receives_default_name(runner, tmp_path):
     source = """
     from typing import Annotated
+    from typing import Any
 
     class CustomNode:
         name: str = ""
+        attributes: dict[Any, Any] = {}
 
         def state(self): return None
         def signature(self): return "signature"

tests/test_data_catalog.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import hashlib
+import pickle
 import sys
 import textwrap
 from pathlib import Path
 
 import pytest
 
+from _pytask.data_catalog_utils import DATA_CATALOG_NAME_FIELD
 from pytask import DataCatalog
 from pytask import ExitCode
 from pytask import PathNode
@@ -198,6 +201,25 @@ def test_adding_a_python_node():
     assert isinstance(data_catalog["node"], PythonNode)
 
 
+def test_reloading_data_catalog_preserves_node_attributes(tmp_path):
+    data_catalog = DataCatalog(_instance_path=tmp_path)
+    _ = data_catalog["node"]
+    assert data_catalog.path is not None
+
+    filename = hashlib.sha256(b"node").hexdigest()
+    path_to_node = data_catalog.path / f"{filename}-node.pkl"
+
+    node = pickle.loads(path_to_node.read_bytes())  # noqa: S301
+    node.attributes["custom"] = "value"
+    path_to_node.write_bytes(pickle.dumps(node))
+
+    reloaded_data_catalog = DataCatalog(_instance_path=tmp_path)
+    reloaded_node = reloaded_data_catalog["node"]
+
+    assert reloaded_node.attributes["custom"] == "value"
+    assert reloaded_node.attributes[DATA_CATALOG_NAME_FIELD] == "default"
+
+
 def test_use_data_catalog_with_provisional_node(runner, tmp_path):
     source = """
     from pathlib import Path

tests/test_execute.py

@@ -421,13 +421,15 @@ def test_custom_node_as_product(runner, tmp_path, product_def, return_def):
     from typing import Any
     from typing import Annotated
     from dataclasses import dataclass
+    from dataclasses import field
     from pytask import Product
 
     @dataclass
     class PickleNode:
         path: Path
         name: str = ""
         signature: str = "id"
+        attributes: dict[Any, Any] = field(default_factory=dict)
 
         def state(self) -> str | None:
             if self.path.exists():
@@ -751,13 +753,15 @@ def test_errors_during_loading_nodes_have_info(runner, tmp_path):
     from pathlib import Path
     from typing import Any
     from dataclasses import dataclass
+    from dataclasses import field
     import pickle
 
     @dataclass
     class PickleNode:
         name: str
         path: Path
         signature: str = "id"
+        attributes: dict[Any, Any] = field(default_factory=dict)
 
         def state(self) -> str | None:
             if self.path.exists():

tests/test_node_protocols.py

@@ -92,12 +92,10 @@ def task_example(
     assert tmp_path.joinpath("out.txt").read_text() == "text"
 
 
-def test_node_protocol_for_custom_nodes_adding_attributes(runner, tmp_path):
+def test_node_protocol_for_custom_nodes_requires_attributes(runner, tmp_path):
     source = """
     from typing import Annotated
-    from pytask import Product
     from dataclasses import dataclass
-    from pathlib import Path
 
     @dataclass
     class CustomNode:
@@ -114,15 +112,11 @@ def load(self, is_product):
         def save(self, value):
             self.value = value
 
-    def task_example(
-        data = CustomNode("custom", "text"),
-        out: Annotated[Path, Product] = Path("out.txt"),
-    ) -> None:
-        out.write_text(data)
+    def task_example() -> Annotated[str, CustomNode("custom", "text")]:
+        return "text"
     """
     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("out.txt").read_text() == "text"
-    assert "FutureWarning" in result.output
+    assert result.exit_code == ExitCode.COLLECTION_FAILED
+    assert "does not follow the 'pytask.PNode' protocol" in result.output