tests + docstrings

tests/test_graph_group.py

@@ -38,6 +38,9 @@ def test_grp():
     assert grp_name in graph._groups
     assert graph.get_group(grp_name) == [node.uuid]
 
+    assert len(graph._groups) == 1
+    assert len(graph) == 1
+
 
 def test_muliple_grps():
     graph = znflow.DiGraph()
@@ -75,6 +78,9 @@ def test_muliple_grps():
     assert graph.get_group(grp_name) == [node.uuid]
     assert graph.get_group(grp_name2) == [node2.uuid]
 
+    assert len(graph._groups) == 2
+    assert len(graph) == 2
+
 
 def test_nested_grps():
     graph = znflow.DiGraph()
@@ -111,6 +117,9 @@ def test_grp_with_existing_nodes():
 
     assert graph.get_group(grp_name) == [node2.uuid]
 
+    assert len(graph._groups) == 1
+    assert len(graph) == 2
+
 
 def test_grp_with_multiple_nodes():
     with znflow.DiGraph() as graph:
@@ -142,3 +151,38 @@ def test_grp_with_multiple_nodes():
     assert grp_name in graph._groups
 
     assert graph.get_group(grp_name) == [node3.uuid, node4.uuid]
+
+    assert len(graph._groups) == 1
+    assert len(graph) == 4
+
+def test_reopen_grps():
+    with znflow.DiGraph() as graph:
+        with graph.group("my_grp") as grp_name:
+            assert graph.active_group == grp_name
+
+            node = PlainNode(1)
+
+        with graph.group("my_grp") as grp_name2:
+            assert graph.active_group == grp_name2
+
+            node2 = PlainNode(2)
+
+    assert graph.active_group is None
+
+    graph.run()
+
+    assert grp_name == "my_grp"
+    assert grp_name2 == grp_name
+
+    assert node.value == 2
+    assert node2.value == 3
+
+    assert node.uuid in graph.nodes
+    assert node2.uuid in graph.nodes
+
+    assert grp_name in graph._groups
+
+    assert graph.get_group(grp_name) == [node.uuid, node2.uuid]
+
+    assert len(graph._groups) == 1
+    assert len(graph) == 2

znflow/graph.py

@@ -161,6 +161,31 @@ def write_graph(self, *args):
 
     @contextlib.contextmanager
     def group(self, name: str) -> typing.Generator[str, None, None]:
+        """Create a group of nodes.
+        
+        Allows to group nodes together, independent of their order in the graph.
+        
+        The group can be created within the graph context manager.
+        Alternatively, the group can be created outside of the graph context manager,
+        implicitly opening and closing the graph context manager.
+
+        Attributes
+        ----------
+            name : str
+                Name of the group. If the name is already used, the nodes will be added
+                to the existing group.
+        
+        Raises
+        ------
+            TypeError
+                If a group with the same name is already active. Nested groups are not
+                supported.
+        
+        Yields
+        ------
+            str
+                Name of the group.
+        """
         if self.active_group is not None:
             raise TypeError(
                 f"Nested groups are not supported. Group with name '{self.active_group}'"