Preserve data catalog node attributes on reload

Author: tobiasraabe

src/_pytask/data_catalog.py

@@ -101,7 +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
-            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:

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