BUG: TreeObject.remove_child non-PdfObject assignments

PyPDF2/generic/_data_structures.py

@@ -374,7 +374,7 @@ def addChild(self, child: Any, pdf: Any) -> None:  # pragma: no cover
         deprecate_with_replacement("addChild", "add_child")
         self.add_child(child, pdf)
 
-    def add_child(self, child: Any, pdf: Any) -> None:  # PdfReader
+    def add_child(self, child: Any, pdf: Any) -> None:  # PdfWriter
         child_obj = child.get_object()
         child = pdf.get_reference(child_obj)
         assert isinstance(child, IndirectObject)
@@ -446,13 +446,14 @@ def remove_child(self, child: Any) -> None:
                         next_obj = next_ref.get_object()
                         next_obj[NameObject("/Prev")] = prev_ref
                         prev[NameObject("/Next")] = next_ref
-                        self[NameObject("/Count")] -= 1
                     else:
                         # Removing last tree node
                         assert cur == last
                         del prev[NameObject("/Next")]
                         self[NameObject("/Last")] = prev_ref
-                        self[NameObject("/Count")] -= 1
+                    self[NameObject("/Count")] = NumberObject(
+                        self[NameObject("/Count")] - 1
+                    )
                 found = True
                 break
 

tests/test_generic.py

@@ -435,6 +435,33 @@ def get_object(self):
     assert exc.value.args[0] == "Removal couldn't find item in tree"
 
 
+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
+
+    tree = TreeObject()
+    child = DictionaryObject()
+    writer._add_object(tree)
+
+    child_ref = writer._add_object(child)
+    tree.add_child(child_ref, writer)
+    assert tree[NameObject("/Count")] == 1
+
+    child2 = TreeObject()
+    child2_ref = writer._add_object(child2)
+    tree.add_child(child2_ref, writer)
+    assert tree[NameObject("/Count")] == 2
+    tree.remove_child(child2)
+
+
 def test_remove_child_in_tree():
     pdf = RESOURCE_ROOT / "form.pdf"