move group updates (#106)

* move group updates

* add tests

* [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>

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "znflow"
-version = "0.2.0"
+version = "0.2.1"
 description = "A general purpose framework for building and running computational graphs."
 authors = ["zincwarecode <[email protected]>"]
 license = "Apache-2.0"

tests/test_graph_build_error.py

@@ -0,0 +1,40 @@
+"""Test for raising an error when building a graph."""
+
+import dataclasses
+
+import znflow
+
+
+@dataclasses.dataclass
+class MyNode(znflow.Node):
+    value: int
+
+    def run(self):
+        pass
+
+
+def test_graph_build_exception():
+    graph = znflow.DiGraph()
+
+    try:
+        with graph:
+            node = MyNode(value=42)
+            raise ValueError("This is a test")
+    except ValueError:
+        pass
+
+    assert node.uuid in graph
+
+
+def test_group_build_exception():
+    graph = znflow.DiGraph()
+
+    try:
+        with graph.group("group") as grp:
+            node = MyNode(value=42)
+            raise ValueError("This is a test")
+    except ValueError:
+        pass
+
+    assert node.uuid in graph
+    assert node.uuid in grp

tests/test_graph_group.py

@@ -35,10 +35,10 @@ def test_grp():
     assert grp.names == ("my_grp",)
     assert node.value == 2
     assert node.uuid in graph.nodes
-    assert grp.names in graph._groups
+    assert grp.names in graph.groups
     assert graph.get_group("my_grp").uuids == [node.uuid]
 
-    assert len(graph._groups) == 1
+    assert len(graph.groups) == 1
     assert len(graph) == 1
 
 
@@ -72,13 +72,13 @@ def test_muliple_grps():
     assert node.uuid in graph.nodes
     assert node2.uuid in graph.nodes
 
-    assert grp.names in graph._groups
-    assert grp2.names in graph._groups
+    assert grp.names in graph.groups
+    assert grp2.names in graph.groups
 
     assert graph.get_group(*grp.names).uuids == [node.uuid]
     assert graph.get_group(*grp2.names).uuids == [node2.uuid]
 
-    assert len(graph._groups) == 2
+    assert len(graph.groups) == 2
     assert len(graph) == 2
 
 
@@ -113,11 +113,11 @@ def test_grp_with_existing_nodes():
     assert node.uuid in graph.nodes
     assert node2.uuid in graph.nodes
 
-    assert grp.names in graph._groups
+    assert grp.names in graph.groups
 
     assert graph.get_group(*grp.names).uuids == [node2.uuid]
 
-    assert len(graph._groups) == 1
+    assert len(graph.groups) == 1
     assert len(graph) == 2
 
 
@@ -148,11 +148,11 @@ def test_grp_with_multiple_nodes():
     assert node3.uuid in graph.nodes
     assert node4.uuid in graph.nodes
 
-    assert grp.names in graph._groups
+    assert grp.names in graph.groups
 
     assert graph.get_group(*grp.names).uuids == [node3.uuid, node4.uuid]
 
-    assert len(graph._groups) == 1
+    assert len(graph.groups) == 1
     assert len(graph) == 4
 
 
@@ -181,11 +181,11 @@ def test_reopen_grps():
     assert node.uuid in graph.nodes
     assert node2.uuid in graph.nodes
 
-    assert grp.names in graph._groups
+    assert grp.names in graph.groups
 
     assert graph.get_group(*grp.names).uuids == [node.uuid, node2.uuid]
 
-    assert len(graph._groups) == 1
+    assert len(graph.groups) == 1
     assert len(graph) == 2
 
 
@@ -204,7 +204,7 @@ def test_tuple_grp_names():
     assert grp.names == ("grp", "1")
     assert node.value == 2
     assert node.uuid in graph.nodes
-    assert grp.names in graph._groups
+    assert grp.names in graph.groups
     assert graph.get_group(*grp.names).uuids == [node.uuid]
 
 

tests/test_znflow.py

@@ -5,4 +5,4 @@
 
 def test_version():
     """Test the version."""
-    assert znflow.__version__ == "0.2.0"
+    assert znflow.__version__ == "0.2.1"

znflow/graph.py

@@ -59,7 +59,7 @@ def __init__(
         """
         self.disable = disable
         self.immutable_nodes = immutable_nodes
-        self._groups = {}
+        self.groups = {}
         self.active_group: typing.Union[Group, None] = None
         self.deployment = deployment or VanillaDeployment()
         self.deployment.set_graph(self)
@@ -244,21 +244,29 @@ def group(self, *names: str) -> typing.Generator[Group, None, None]:
 
         existing_nodes = self.get_sorted_nodes()
 
-        group = self._groups.get(names, Group(names=names, uuids=[], graph=self))
+        group = self.groups.get(names, Group(names=names, uuids=[], graph=self))
+
+        def finalize_group():
+            for node_uuid in self.nodes:
+                if node_uuid not in existing_nodes:
+                    self.groups[group.names] = group
+                    group.uuids.append(node_uuid)
 
         try:
             self.active_group = group
             if get_graph() is empty_graph:
                 with self:
-                    yield group
+                    try:
+                        yield group
+                    finally:
+                        finalize_group()
             else:
-                yield group
+                try:
+                    yield group
+                finally:
+                    finalize_group()
         finally:
             self.active_group = None
-            for node_uuid in self.nodes:
-                if node_uuid not in existing_nodes:
-                    self._groups[group.names] = group
-                    group.uuids.append(node_uuid)
 
     def get_group(self, *names: str) -> Group:
-        return self._groups[names]
+        return self.groups[names]