UPDATE: add new methods to remove node with/without its leaves

In `objects.NodeCollection` class, method `remove_node_and_its_leaves()`
is added for removing a node and all its leaf nodes.

And in `object.Node` class, method `remove_all_leaves()` is added
for removing all links of leaf nodes from specific node.

With these new methods, user now can choose to remove all leaf nodes
or not when they are going to remove a node. So that users won't
need to remove those leaf nodes one by one manually.

codememo/components.py

@@ -11,6 +11,7 @@
     NodeCollection,
 )
 from .events import NodeEvent, NodeEventRegistry
+from .exceptions import NodeRemovalException
 from .internal import GlobalState
 
 CODE_CHAR_WIDTH = 8
@@ -1281,6 +1282,42 @@ def remove_node_component(self, node_component):
             if self.id_selected == node_component_id:
                 self.id_selected = -1   # reset index of selected node
                 self.selected_node = None
+        except NodeRemovalException as ex_node_removal:
+            def _remove_node_and_leaves(node_component):
+                uuids = [v.node.uuid for v in self.node_components]
+                removed = self.node_collection.remove_node_and_its_leaves(node_component.node)
+
+                for node in removed:
+                    idx = uuids.index(node.uuid)
+                    self.node_components.pop(idx)
+                    uuids.pop(idx)
+
+                self.links = self.node_collection.resolve_links()
+                self.id_selected = -1
+                self.selected_node = None
+
+            def _remove_node_but_keep_leaves(node_component):
+                node_component.node.remove_all_leaves()
+                self.node_collection.remove_node(node_component.node)
+
+                idx = self.node_components.index(node_component)
+                self.node_components.pop(idx)
+
+                self.links = self.node_collection.resolve_links()
+                if self.id_selected == node_component.id:
+                    self.id_selected = -1
+                    self.selected_node = None
+
+            msg = (
+                'There are still existing leaf nodes, do you want to remove them all?\n' +
+                'Yes: remove all leaves\nNo: keep leaves\nCancel: cancel operation'
+            )
+            self.confirmation_modal = ConfirmationModal(
+                'Confirm', msg,
+                callback_yes=lambda: _remove_node_and_leaves(node_component),
+                callback_no=lambda: _remove_node_but_keep_leaves(node_component),
+                show_cancel_button=True,
+            )
         except Exception as ex:
             GlobalState().push_error(ex)
 

codememo/objects.py

@@ -270,6 +270,12 @@ def remove_leaf_by_index(self, idx):
         node = self.leaves.pop(idx)
         node.reset_root(self)
 
+    def remove_all_leaves(self):
+        """Remove all leaf nodes from this node."""
+        for i in range(len(self.leaves)):
+            current_leaf = self.leaves[0]
+            self.remove_leaf(current_leaf)
+
 
 class NodeLink(object):
     """A link indicates the relation between root and leaf node."""
@@ -382,6 +388,33 @@ def remove_node(self, target):
         for root in target.roots:
             root.remove_leaf(target)
 
+    def remove_node_and_its_leaves(self, target):
+        """Remove node and all its leaves from this collection.
+
+        Parameters
+        ----------
+        target : Node
+            Target node to be removed.
+        """
+        def _remove_node_link(node, removed, visited):
+            if node in visited:
+                return
+            visited.add(node)
+
+            for i in range(len(node.leaves)):
+                current_leaf = node.leaves[0]
+                _remove_node_link(current_leaf, removed, visited)
+
+            for root in node.roots:
+                root.remove_leaf(node)
+            removed.append(node)
+
+        removed, visited = [], set()
+        _remove_node_link(target, removed, visited)
+        for node in removed:
+            self.remove_node(node)
+        return removed
+
     def remove_root_reference(self, target, root):
         """Remove root reference (a.k.a. root node) of given node.
 

tests/test_objects.py

@@ -74,11 +74,11 @@ def dummy_nodes_multiple_trees():
         ('0_2', 'def buzz():\n    print("buzz")', 'python'),
         ('0_3', 'def gin():\n    print("gin")', 'python'),
         ('1_0', 'def fizz():\n    print("fizz")', 'python'),
-        ('2_1', 'def foo():\n    print("foo")', 'python'),
-        ('2_2', 'def bar():\n    print("bar")', 'python'),
-        ('2_3', 'def buzz():\n    print("buzz")', 'python'),
-        ('2_4', 'def gin():\n    print("gin")', 'python'),
-        ('2_5', 'def fizz():\n    print("fizz")', 'python'),
+        ('2_0', 'def foo():\n    print("foo")', 'python'),
+        ('2_1', 'def bar():\n    print("bar")', 'python'),
+        ('2_2', 'def buzz():\n    print("buzz")', 'python'),
+        ('2_3', 'def gin():\n    print("gin")', 'python'),
+        ('2_4', 'def fizz():\n    print("fizz")', 'python'),
     ]
     snippets = [Snippet(v[0], v[1], lang=v[2]) for v in data]
     nodes = [Node(v) for v in snippets]
@@ -197,22 +197,22 @@ def test__add_leaf__self_reference_duplicate(self, dummy_nodes):
             A.add_leaf(A)
 
     def test__add_leaf__multiple_roots(self, dummy_nodes):
-        A, B, C = dummy_nodes[0], dummy_nodes[1], dummy_nodes[2]
+        A, B, C = dummy_nodes[:3]
         A.add_leaf(B)
         C.add_leaf(B)
         assert A in B.roots
         assert C in B.roots
 
     def test__add_leaf__exceed_range(self, dummy_nodes):
-        A, B = dummy_nodes[0], dummy_nodes[1]
+        A, B = dummy_nodes[:2]
         n_lines = A.snippet.n_lines
         with pytest.raises(ValueError, match='should be in the range'):
             A.add_leaf(B, 0)
         with pytest.raises(ValueError, match='should be in the range'):
             A.add_leaf(B, n_lines + 1)
 
     def test__remove_leaf(self, dummy_nodes):
-        A, B, C = dummy_nodes[0], dummy_nodes[1], dummy_nodes[2]
+        A, B, C = dummy_nodes[:3]
         A.add_leaf(B)
         A.add_leaf(C)
 
@@ -225,6 +225,65 @@ def test__remove_leaf(self, dummy_nodes):
         with pytest.raises(NodeRemovalException, match='not a leaf'):
             A.remove_leaf(B)
 
+    def test__remove_all_leaves(self, dummy_nodes):
+        A, B, C, D = dummy_nodes[:4]
+        A.add_leaf(B)
+        A.add_leaf(C)
+        B.add_leaf(D)
+
+        # Current graph:
+        # A --> B --> D
+        #   \-> C
+        assert all([leaf in A.leaves for leaf in [B, C]])
+        assert D in B.leaves
+
+        A.remove_all_leaves()
+
+        # Current graph:
+        # A     B --> D
+        #       C
+        assert A.leaves == []
+        assert B.roots == [] and B.leaves == [D]
+        assert C.roots == []
+
+    def test__remove_all_leaves__circular_references(self, dummy_nodes):
+        A, B, C, D = dummy_nodes[:4]
+        A.add_leaf(B)
+        B.add_leaf(C)
+        C.add_leaf(D)
+        D.add_leaf(A)
+
+        # Current graph:
+        # -> A --> B --> C --> D -
+        # |----------------------|
+        assert A.leaves == [B] and B.roots == [A]
+        assert B.leaves == [C] and C.roots == [B]
+        assert C.leaves == [D] and D.roots == [C]
+        assert D.leaves == [A] and A.roots == [D]
+
+        # Current graph:
+        # B --> C --> D --> A
+        A.remove_all_leaves()
+        assert B.leaves == [C] and C.roots == [B]
+        assert C.leaves == [D] and D.roots == [C]
+        assert D.leaves == [A] and A.roots == [D]
+        assert A.leaves == []
+
+    def test__remove_all_leaves__self_reference(self, dummy_nodes):
+        A = dummy_nodes[0]
+        A.add_leaf(A)
+
+        # Current graph:
+        # -> A -
+        # |----|
+        assert A.leaves == [A]
+        assert A.roots == [A]
+
+        # Current graph:
+        #    A
+        A.remove_all_leaves()
+        assert A.leaves == []
+        assert A.roots == []
 
 class TestNodeCollection:
     def test__add_leaf_reference(self, dummy_nodes):
@@ -275,6 +334,68 @@ def test__remove_node(self, dummy_nodes_multiple_trees):
         with pytest.raises(NodeRemovalException, match='there are remaining leaves'):
             node_collection.remove_node(root)
 
+    def test__remove_node_and_its_leaves__multiple_trees(
+        self, dummy_nodes_multiple_trees
+    ):
+        nodes, desired_links, _ = dummy_nodes_multiple_trees
+        node_collection = NodeCollection(nodes)
+
+        node_0_0_to_0_3 = nodes[:4]
+        node_0_0 = nodes[0]
+        links_to_be_removed, remaining_links = desired_links[:3], desired_links[3:]
+        removed = node_collection.remove_node_and_its_leaves(node_0_0)
+        assert set(removed) == set(node_0_0_to_0_3)
+
+        links = node_collection.resolve_links()
+        assert all([removed_link not in links for removed_link in links_to_be_removed])
+        assert len(links) == len(remaining_links)
+        assert all([remaining_link in links for remaining_link in remaining_links])
+
+        node_2_1_to_2_4 = nodes[-4:]
+        node_2_1 = nodes[-4]
+        links_to_be_removed, remaining_links = desired_links[-3:], []
+        removed = node_collection.remove_node_and_its_leaves(node_2_1)
+        assert set(removed) == set(node_2_1_to_2_4)
+
+        # `node_2_1` is a leaf node of `node_2_0`, so the link <2_0, 2_1> will
+        # also be removed after `node_2_1` is removed. Hence that there is no
+        # remaining links now.
+        links = node_collection.resolve_links()
+        assert links == []
+
+    def test__remove_node_and_its_leaves__circular_references(
+        self, dummy_nodes_circular_references
+    ):
+        nodes, desired_links, _ = dummy_nodes_circular_references
+        node_collection = NodeCollection(nodes)
+
+        node_1_0_to_1_2 = nodes[1:4]
+        node_1_0 = nodes[1]
+        links_to_be_removed, remaining_links = desired_links[:3], desired_links[3:]
+        removed = node_collection.remove_node_and_its_leaves(node_1_0)
+        assert set(removed) == set(node_1_0_to_1_2)
+
+        links = node_collection.resolve_links()
+        assert all([removed_link not in links for removed_link in links_to_be_removed])
+        assert len(links) == len(remaining_links)
+        assert all([remaining_link in links for remaining_link in remaining_links])
+
+    def test__remove_node_and_its_leaves__self_reference(self, dummy_nodes):
+        nodes = dummy_nodes
+        A = nodes[0]
+        A.add_leaf(A)
+        assert A.leaves == [A] and A.roots == [A]
+
+        node_collection = NodeCollection([A])
+        links = node_collection.resolve_links()
+        assert links == [NodeLink(A, 0, A, 0)]
+
+        node_collection.remove_node_and_its_leaves(A)
+        assert len(node_collection) == 0
+
+        links = node_collection.resolve_links()
+        assert links == []
+
     def test__to_dict(self, dummy_node_collection_data):
         data = dummy_node_collection_data
         node_collection = NodeCollection.from_dict(data)