check nested dataclasses works with list dataclass as parameter (#939)

* test that should pass

* bugfix nested dataclasses

* bump version

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: PythonFZ

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "zntrack"
-version = "0.8.9"
+version = "0.8.10"
 description = "Create, Run and Benchmark DVC Pipelines in Python"
 authors = [
     { name = "Fabian Zills", email = "fzills@icp.uni-stuttgart.de" },

tests/integration/test_nested_dataclass_deps.py

@@ -56,3 +56,160 @@ def test_nested_dc_deps(proj_path):
     project.repro()
 
     assert md.from_rev().result == "Berendsen thermostat '1.0'"
+
+
+# test nested lists
+
+
+@dataclasses.dataclass(frozen=True)
+class FuncOne:
+    value: int = 1
+
+    def get_value(self):
+        return self.value
+
+
+@dataclasses.dataclass
+class FuncCollector:
+    funcs: list[FuncOne] = dataclasses.field(default_factory=list)
+
+    def get_values(self):
+        return [func.get_value() for func in self.funcs]
+
+
+class FuncNode(zntrack.Node):
+    collector: FuncCollector = zntrack.deps()
+    result: list[int] = zntrack.outs()
+
+    def run(self):
+        self.result = self.collector.get_values()
+
+
+def test_nested_list_deps(proj_path):
+    project = zntrack.Project()
+    collector = FuncCollector(funcs=[FuncOne(value=1), FuncOne(value=2)])
+
+    with project:
+        node = FuncNode(collector=collector)
+
+    project.repro()
+    assert node.from_rev().result == [1, 2]
+
+    collector.funcs.append(FuncOne(value=3))
+    project.repro()
+
+    assert node.from_rev().result == [1, 2, 3]
+
+
+# Additional test classes for other collection types
+@dataclasses.dataclass
+class FuncTupleCollector:
+    funcs: tuple[FuncOne, ...] = dataclasses.field(default_factory=tuple)
+
+    def get_values(self):
+        return [func.get_value() for func in self.funcs]
+
+
+@dataclasses.dataclass
+class FuncSetCollector:
+    funcs: set[FuncOne] = dataclasses.field(default_factory=set)
+
+    def get_values(self):
+        # Sort for deterministic testing
+        return sorted([func.get_value() for func in self.funcs])
+
+
+@dataclasses.dataclass
+class FuncDictCollector:
+    funcs: dict[str, FuncOne] = dataclasses.field(default_factory=dict)
+
+    def get_values(self):
+        # Return sorted list of values for deterministic testing
+        return sorted([func.get_value() for func in self.funcs.values()])
+
+
+class FuncTupleNode(zntrack.Node):
+    collector: FuncTupleCollector = zntrack.deps()
+    result: list[int] = zntrack.outs()
+
+    def run(self):
+        self.result = self.collector.get_values()
+
+
+class FuncSetNode(zntrack.Node):
+    collector: FuncSetCollector = zntrack.deps()
+    result: list[int] = zntrack.outs()
+
+    def run(self):
+        self.result = self.collector.get_values()
+
+
+class FuncDictNode(zntrack.Node):
+    collector: FuncDictCollector = zntrack.deps()
+    result: list[int] = zntrack.outs()
+
+    def run(self):
+        self.result = self.collector.get_values()
+
+
+def test_nested_tuple_deps(proj_path):
+    project = zntrack.Project()
+    collector = FuncTupleCollector(funcs=(FuncOne(value=1), FuncOne(value=2)))
+
+    with project:
+        node = FuncTupleNode(collector=collector)
+
+    project.repro()
+    assert node.from_rev().result == [1, 2]
+
+    # Test with updated tuple - need to update the same node instance
+    node.collector = FuncTupleCollector(
+        funcs=(FuncOne(value=1), FuncOne(value=2), FuncOne(value=3))
+    )
+    project.repro()
+
+    assert node.from_rev().result == [1, 2, 3]
+
+
+def test_nested_set_deps(proj_path):
+    project = zntrack.Project()
+    collector = FuncSetCollector(funcs={FuncOne(value=1), FuncOne(value=2)})
+
+    with project:
+        node = FuncSetNode(collector=collector)
+
+    project.repro()
+    assert node.from_rev().result == [1, 2]
+
+    # Test with updated set - need to update the same node instance
+    node.collector = FuncSetCollector(
+        funcs={FuncOne(value=1), FuncOne(value=2), FuncOne(value=3)}
+    )
+    project.repro()
+
+    assert node.from_rev().result == [1, 2, 3]
+
+
+def test_nested_dict_deps(proj_path):
+    project = zntrack.Project()
+    collector = FuncDictCollector(
+        funcs={"first": FuncOne(value=1), "second": FuncOne(value=2)}
+    )
+
+    with project:
+        node = FuncDictNode(collector=collector)
+
+    project.repro()
+    assert node.from_rev().result == [1, 2]
+
+    # Test with updated dict - need to update the same node instance
+    node.collector = FuncDictCollector(
+        funcs={
+            "first": FuncOne(value=1),
+            "second": FuncOne(value=2),
+            "third": FuncOne(value=3),
+        }
+    )
+    project.repro()
+
+    assert node.from_rev().result == [1, 2, 3]

uv.lock

@@ -23,6 +23,42 @@
 from .utils import module_handler
 
 
+def _reconstruct_value_recursively(value):
+    """Recursively reconstruct values, handling dataclasses in collections."""
+    if isinstance(value, dict) and "_cls" in value:
+        # This is a nested dataclass, reconstruct it
+        cls_path = value["_cls"]
+        if not isinstance(cls_path, str) or "." not in cls_path:
+            raise ValueError(f"Invalid class path format: {cls_path}")
+        try:
+            module_name, class_name = cls_path.rsplit(".", 1)
+            module = importlib.import_module(module_name)
+            nested_cls = getattr(module, class_name)
+        except (ImportError, AttributeError, ValueError) as e:
+            raise ImportError(f"Failed to import class {cls_path}: {e}") from e
+
+        if not dataclasses.is_dataclass(nested_cls):
+            raise TypeError(f"Class {cls_path} is not a dataclass")
+
+        # Recursively process nested parameters
+        nested_params = _reconstruct_nested_dataclasses(value)
+        try:
+            return nested_cls(**nested_params)
+        except TypeError as e:
+            raise TypeError(f"Failed to instantiate {cls_path}: {e}") from e
+    elif isinstance(value, list):
+        return [_reconstruct_value_recursively(item) for item in value]
+    elif isinstance(value, tuple):
+        return tuple(_reconstruct_value_recursively(item) for item in value)
+    elif isinstance(value, set):
+        # Note: Sets are serialized as lists, so this won't be called during normal flow
+        return {_reconstruct_value_recursively(item) for item in value}
+    elif isinstance(value, dict):
+        return {k: _reconstruct_value_recursively(v) for k, v in value.items()}
+    else:
+        return value
+
+
 def _reconstruct_nested_dataclasses(params: dict) -> dict:
     """Recursively reconstruct nested dataclasses from their dictionary representation."""
     if not isinstance(params, dict):
@@ -32,29 +68,7 @@ def _reconstruct_nested_dataclasses(params: dict) -> dict:
     for key, value in params.items():
         if key == "_cls":
             continue
-        if isinstance(value, dict) and "_cls" in value:
-            # This is a nested dataclass, reconstruct it
-            cls_path = value["_cls"]
-            if not isinstance(cls_path, str) or "." not in cls_path:
-                raise ValueError(f"Invalid class path format: {cls_path}")
-            try:
-                module_name, class_name = cls_path.rsplit(".", 1)
-                module = importlib.import_module(module_name)
-                nested_cls = getattr(module, class_name)
-            except (ImportError, AttributeError, ValueError) as e:
-                raise ImportError(f"Failed to import class {cls_path}: {e}") from e
-
-            if not dataclasses.is_dataclass(nested_cls):
-                raise TypeError(f"Class {cls_path} is not a dataclass")
-
-            # Recursively process nested parameters
-            nested_params = _reconstruct_nested_dataclasses(value)
-            try:
-                reconstructed_params[key] = nested_cls(**nested_params)
-            except TypeError as e:
-                raise TypeError(f"Failed to instantiate {cls_path}: {e}") from e
-        else:
-            reconstructed_params[key] = value
+        reconstructed_params[key] = _reconstruct_value_recursively(value)
     return reconstructed_params
 
 

zntrack/converter.py

@@ -11,6 +11,25 @@
 from zntrack.utils import module_handler
 
 
+def _convert_value_recursively(value):
+    """Recursively convert values, handling dataclasses in collections."""
+    if dataclasses.is_dataclass(value) and not isinstance(
+        value, (Node, znflow.Connection, znflow.CombinedConnections)
+    ):
+        return _dataclass_to_dict(value)
+    elif isinstance(value, list):
+        return [_convert_value_recursively(item) for item in value]
+    elif isinstance(value, tuple):
+        return tuple(_convert_value_recursively(item) for item in value)
+    elif isinstance(value, set):
+        # Sets cannot contain dicts (unhashable), so convert to list
+        return [_convert_value_recursively(item) for item in value]
+    elif isinstance(value, dict):
+        return {k: _convert_value_recursively(v) for k, v in value.items()}
+    else:
+        return copy.deepcopy(value)
+
+
 def _dataclass_to_dict(object) -> dict:
     """Convert a dataclass to a dictionary excluding certain keys."""
     exclude_fields = []
@@ -35,13 +54,7 @@ def _dataclass_to_dict(object) -> dict:
             continue
 
         value = getattr(object, field.name)
-        if dataclasses.is_dataclass(value) and not isinstance(
-            value, (Node, znflow.Connection, znflow.CombinedConnections)
-        ):
-            # Recursively convert nested dataclasses
-            dc_params[field.name] = _dataclass_to_dict(value)
-        else:
-            dc_params[field.name] = copy.deepcopy(value)
+        dc_params[field.name] = _convert_value_recursively(value)
 
     dc_params["_cls"] = f"{module_handler(object)}.{object.__class__.__name__}"
     return dc_params
@@ -51,24 +64,48 @@ def deps_to_params(self, field):
     if getattr(self.node, field.name) is None:
         return
     content = getattr(self.node, field.name)
-    if isinstance(content, (list, tuple, dict)):
-        new_content = []
-        for val in content if isinstance(content, (list, tuple)) else content.values():
-            if dataclasses.is_dataclass(val) and not isinstance(
-                val, (Node, znflow.Connection, znflow.CombinedConnections)
-            ):
-                # We save the values of the passed dataclasses
-                #  to the params.yaml file to be later used
-                #  by the DataclassContainer to recreate the
-                #  instance with the correct parameters.
-                new_content.append(_dataclass_to_dict(val))
-            elif isinstance(val, (znflow.Connection, znflow.CombinedConnections)):
-                pass
-            else:
-                raise ValueError(
-                    f"Found unsupported type '{type(val)}' ({val}) for DEPS"
-                    f" field '{field.name}' in list"
-                )
+    if isinstance(content, (list, tuple, set, dict)):
+        if isinstance(content, dict):
+            # For dicts, we need to convert both keys and values
+            new_content = {}
+            for key, val in content.items():
+                if dataclasses.is_dataclass(val) and not isinstance(
+                    val, (Node, znflow.Connection, znflow.CombinedConnections)
+                ):
+                    new_content[key] = _dataclass_to_dict(val)
+                elif isinstance(val, (znflow.Connection, znflow.CombinedConnections)):
+                    pass
+                else:
+                    raise ValueError(
+                        f"Found unsupported type '{type(val)}' ({val}) for DEPS"
+                        f" field '{field.name}' in dict"
+                    )
+        else:
+            # For lists, tuples, sets
+            new_content = []
+            for val in content:
+                if dataclasses.is_dataclass(val) and not isinstance(
+                    val, (Node, znflow.Connection, znflow.CombinedConnections)
+                ):
+                    # We save the values of the passed dataclasses
+                    #  to the params.yaml file to be later used
+                    #  by the DataclassContainer to recreate the
+                    #  instance with the correct parameters.
+                    new_content.append(_dataclass_to_dict(val))
+                elif isinstance(val, (znflow.Connection, znflow.CombinedConnections)):
+                    pass
+                else:
+                    raise ValueError(
+                        f"Found unsupported type '{type(val)}' ({val}) for DEPS"
+                        f" field '{field.name}' in collection"
+                    )
+            # Preserve the original collection type
+            if isinstance(content, tuple):
+                new_content = tuple(new_content)
+            elif isinstance(content, set):
+                # Sets cannot contain dicts (unhashable), so convert to list
+                new_content = list(new_content)
+
         if len(new_content) > 0:
             return new_content
     elif dataclasses.is_dataclass(content) and not isinstance(