BUG: TreeObject.remove_child had an assignment issue for Count (#1234)

PyPDF2/generic/_data_structures.py

@@ -406,6 +406,41 @@ def removeChild(self, child: Any) -> None:  # pragma: no cover
         deprecate_with_replacement("removeChild", "remove_child")
         self.remove_child(child)
 
+    def _remove_node_from_tree(
+        self, prev: Any, prev_ref: Any, cur: Any, last: Any
+    ) -> None:
+        """Adjust the pointers of the linked list and tree node count."""
+        next_ref = cur.get(NameObject("/Next"), None)
+        if prev is None:
+            if next_ref:
+                # Removing first tree node
+                next_obj = next_ref.get_object()
+                del next_obj[NameObject("/Prev")]
+                self[NameObject("/First")] = next_ref
+                self[NameObject("/Count")] = NumberObject(
+                    self[NameObject("/Count")] - 1  # type: ignore
+                )
+
+            else:
+                # Removing only tree node
+                assert self[NameObject("/Count")] == 1
+                del self[NameObject("/Count")]
+                del self[NameObject("/First")]
+                if NameObject("/Last") in self:
+                    del self[NameObject("/Last")]
+        else:
+            if next_ref:
+                # Removing middle tree node
+                next_obj = next_ref.get_object()
+                next_obj[NameObject("/Prev")] = prev_ref
+                prev[NameObject("/Next")] = next_ref
+            else:
+                # Removing last tree node
+                assert cur == last
+                del prev[NameObject("/Next")]
+                self[NameObject("/Last")] = prev_ref
+            self[NameObject("/Count")] = NumberObject(self[NameObject("/Count")] - 1)  # type: ignore
+
     def remove_child(self, child: Any) -> None:
         child_obj = child.get_object()
 
@@ -423,40 +458,11 @@ def remove_child(self, child: Any) -> None:
         last = last_ref.get_object()
         while cur is not None:
             if cur == child_obj:
-                if prev is None:
-                    if NameObject("/Next") in cur:
-                        # Removing first tree node
-                        next_ref = cur[NameObject("/Next")]
-                        next_obj = next_ref.get_object()
-                        del next_obj[NameObject("/Prev")]
-                        self[NameObject("/First")] = next_ref
-                        self[NameObject("/Count")] -= 1  # type: ignore
-
-                    else:
-                        # Removing only tree node
-                        assert self[NameObject("/Count")] == 1
-                        del self[NameObject("/Count")]
-                        del self[NameObject("/First")]
-                        if NameObject("/Last") in self:
-                            del self[NameObject("/Last")]
-                else:
-                    if NameObject("/Next") in cur:
-                        # Removing middle tree node
-                        next_ref = cur[NameObject("/Next")]
-                        next_obj = next_ref.get_object()
-                        next_obj[NameObject("/Prev")] = prev_ref
-                        prev[NameObject("/Next")] = next_ref
-                    else:
-                        # Removing last tree node
-                        assert cur == last
-                        del prev[NameObject("/Next")]
-                        self[NameObject("/Last")] = prev_ref
-                    self[NameObject("/Count")] = NumberObject(
-                        self[NameObject("/Count")] - 1
-                    )
+                self._remove_node_from_tree(prev, prev_ref, cur, last)
                 found = True
                 break
 
+            # Go to the next node
             prev_ref = cur_ref
             prev = cur
             if NameObject("/Next") in cur:
@@ -469,11 +475,7 @@ def remove_child(self, child: Any) -> None:
         if not found:
             raise ValueError("Removal couldn't find item in tree")
 
-        del child_obj[NameObject("/Parent")]
-        if NameObject("/Next") in child_obj:
-            del child_obj[NameObject("/Next")]
-        if NameObject("/Prev") in child_obj:
-            del child_obj[NameObject("/Prev")]
+        _reset_node_tree_relationship(child_obj)
 
     def emptyTree(self) -> None:  # pragma: no cover
         deprecate_with_replacement("emptyTree", "empty_tree", "4.0.0")
@@ -496,6 +498,19 @@ def empty_tree(self) -> None:
             del self[NameObject("/Last")]
 
 
+def _reset_node_tree_relationship(child_obj: Any) -> None:
+    """
+    Call this after a node has been removed from a tree.
+
+    This resets the nodes attributes in respect to that tree.
+    """
+    del child_obj[NameObject("/Parent")]
+    if NameObject("/Next") in child_obj:
+        del child_obj[NameObject("/Next")]
+    if NameObject("/Prev") in child_obj:
+        del child_obj[NameObject("/Prev")]
+
+
 class StreamObject(DictionaryObject):
     def __init__(self) -> None:
         self.__data: Optional[str] = None

tests/test_generic.py

@@ -438,28 +438,40 @@ def get_object(self):
 def test_remove_child_found_in_tree():
     writer = PdfWriter()
 
-    class ChildDummy:
-        def __init__(self, parent):
-            self.parent = parent
-
-        def get_object(self):
-            tree = DictionaryObject()
-            tree[NameObject("/Parent")] = self.parent
-            return tree
-
+    # Add Tree
     tree = TreeObject()
-    child = DictionaryObject()
     writer._add_object(tree)
 
-    child_ref = writer._add_object(child)
-    tree.add_child(child_ref, writer)
+    # Add first child
+    # It's important to set a value, otherwise the writer.get_reference will
+    # return the same object when a second child is added.
+    child1 = TreeObject()
+    child1[NameObject("/Foo")] = TextStringObject("bar")
+    child1_ref = writer._add_object(child1)
+    tree.add_child(child1_ref, writer)
     assert tree[NameObject("/Count")] == 1
 
+    # Add second child
     child2 = TreeObject()
+    child2[NameObject("/Foo")] = TextStringObject("baz")
     child2_ref = writer._add_object(child2)
     tree.add_child(child2_ref, writer)
     assert tree[NameObject("/Count")] == 2
+
+    # Remove child
     tree.remove_child(child2)
+    assert tree[NameObject("/Count")] == 1
+
+    # Add new child
+    child3 = TreeObject()
+    child3_ref = writer._add_object(child3)
+    tree.add_child(child3_ref, writer)
+    assert tree[NameObject("/Count")] == 2
+
+    # Remove child
+    child1 = tree[NameObject("/First")]
+    tree.remove_child(child1)
+    assert tree[NameObject("/Count")] == 1
 
 
 def test_remove_child_in_tree():