Update znflow (#548)

* add eager test

* first try in supporting AddedConnections

* convert tuple to list

* use znflow from pypi

* add test for combined dict

* add error message

* add comments

* use dict comprehension

* add 'name' to '_protected_'

Author: PythonFZ

poetry.lock

@@ -18,8 +18,8 @@ typer = "^0.7.0"
 
 dot4dict = "^0.1.1"
 zninit = "^0.1.9"
-znflow = "^0.1.5"
 znjson = "^0.2.2"
+znflow = "^0.1.6"
 
 
 [tool.poetry.urls]

pyproject.toml

@@ -0,0 +1,59 @@
+import pytest
+
+import zntrack
+
+
+class GenerateList(zntrack.Node):
+    size = zntrack.zn.params(10)
+    outs = zntrack.zn.outs()
+
+    def run(self):
+        self.outs = list(range(self.size))
+
+
+class AddOneToList(zntrack.Node):
+    data = zntrack.zn.deps()
+    outs = zntrack.zn.outs()
+
+    def run(self) -> None:
+        self.outs = [x + 1 for x in self.data]
+
+
+class AddOneToDict(zntrack.Node):
+    data = zntrack.zn.deps()
+    outs = zntrack.zn.outs()
+
+    def run(self) -> None:
+        self.outs = {k: [x + 1 for x in v] for k, v in self.data.items()}
+
+
+@pytest.mark.parametrize("eager", [True, False])
+def test_combine(proj_path, eager):
+    with zntrack.Project() as proj:
+        a = GenerateList(size=1, name="a")
+        b = GenerateList(size=2, name="b")
+        c = GenerateList(size=3, name="c")
+
+        added = AddOneToList(data=a.outs + b.outs + c.outs)
+
+    proj.run(eager=eager)
+    if not eager:
+        added.load()
+
+    assert added.outs == [1] + [1, 2] + [1, 2, 3]
+
+
+@pytest.mark.parametrize("eager", [True, False])
+def test_combine_dict(proj_path, eager):
+    with zntrack.Project() as proj:
+        a = GenerateList(size=1, name="a")
+        b = GenerateList(size=2, name="b")
+        c = GenerateList(size=3, name="c")
+
+        added = AddOneToDict(data={x.name: x.outs for x in [a, b, c]})
+
+    proj.run(eager=eager)
+    if not eager:
+        added.load()
+
+    assert added.outs == {"a": [1], "b": [1, 2], "c": [1, 2, 3]}

tests/integration/test_combine_lists.py

@@ -1,6 +1,7 @@
-import zntrack
 import znflow
 
+import zntrack
+
 
 class NodeWithProperty(zntrack.Node):
     params = zntrack.zn.params(None)

tests/integration/test_misc.py

@@ -88,6 +88,8 @@ class Node(zninit.ZnInit, znflow.Node):
     name: str = _NameDescriptor(None)
     _name_ = None
 
+    _protected_ = znflow.Node._protected_ + ["name"]
+
     def _post_load_(self) -> None:
         """Post load hook.
 

zntrack/core/node.py

@@ -7,6 +7,7 @@
 
 import pandas as pd
 import yaml
+import znflow
 import znflow.utils
 import zninit
 import znjson
@@ -38,6 +39,36 @@ def decode(self, value: str) -> znflow.Connection:
         return znflow.Connection(**value)
 
 
+class CombinedConnectionsConverter(znjson.ConverterBase):
+    """Convert a znflow.Connection object to dict and back."""
+
+    level = 100
+    representation = "znflow.CombinedConnections"
+    instance = znflow.CombinedConnections
+
+    def encode(self, obj: znflow.CombinedConnections) -> dict:
+        """Convert the znflow.Connection object to dict."""
+        if obj.item is not None:
+            raise NotImplementedError(
+                "znflow.CombinedConnections getitem is not supported yet."
+            )
+        return dataclasses.asdict(obj)
+
+    def decode(self, value: str) -> znflow.CombinedConnections:
+        """Create znflow.Connection object from dict."""
+        connections = []
+        for item in value["connections"]:
+            if isinstance(item, dict):
+                # @nodify functions aren't support as 'zn.deps'
+                # Nodes directly aren't supported because they aren't lists
+                connections.append(znflow.Connection(**item))
+            else:
+                # For the case that item is already a znflow.Connection
+                connections.append(item)
+        value["connections"] = connections
+        return znflow.CombinedConnections(**value)
+
+
 class SliceConverter(znjson.ConverterBase):
     """Convert a znflow.Connection object to dict and back."""
 
@@ -269,9 +300,21 @@ def get_files(self, instance) -> list:
         files = []
 
         value = getattr(instance, self.name)
+        # TODO use IterableHandler?
 
+        if isinstance(value, dict):
+            value = list(value.values())
         if not isinstance(value, (list, tuple)):
             value = [value]
+        if isinstance(value, tuple):
+            value = list(value)
+
+        others = []
+        for node in value:
+            if isinstance(node, znflow.CombinedConnections):
+                others.extend(node.connections)
+
+        value.extend(others)
 
         for node in value:
             if node is None:
@@ -298,7 +341,7 @@ def save(self, instance: "Node"):
             value,
             instance,
             encoder=znjson.ZnEncoder.from_converters(
-                [ConnectionConverter], add_default=True
+                [ConnectionConverter, CombinedConnectionsConverter], add_default=True
             ),
         )
 
@@ -313,7 +356,9 @@ def get_data(self, instance: "Node") -> any:
 
         value = json.loads(
             json.dumps(value),
-            cls=znjson.ZnDecoder.from_converters(ConnectionConverter, add_default=True),
+            cls=znjson.ZnDecoder.from_converters(
+                [ConnectionConverter, CombinedConnectionsConverter], add_default=True
+            ),
         )
 
         # Up until here we have connection objects. Now we need