ENH: Add += and -= operators to ArrayObject (#2510)

Author: pubpub-zz

pypdf/generic/_data_structures.py

@@ -31,6 +31,7 @@
 
 import logging
 import re
+import sys
 from io import BytesIO
 from typing import (
     Any,
@@ -70,6 +71,7 @@
 from ..errors import STREAM_TRUNCATED_PREMATURELY, PdfReadError, PdfStreamError
 from ._base import (
     BooleanObject,
+    ByteStringObject,
     FloatObject,
     IndirectObject,
     NameObject,
@@ -81,6 +83,11 @@
 from ._fit import Fit
 from ._utils import read_hex_string_from_stream, read_string_from_stream
 
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
 logger = logging.getLogger(__name__)
 NumberSigns = b"+-"
 IndirectPattern = re.compile(rb"[+-]?(\d+)\s+(\d+)\s+R[^a-zA-Z]")
@@ -121,6 +128,65 @@ def items(self) -> Iterable[Any]:
         """Emulate DictionaryObject.items for a list (index, object)."""
         return enumerate(self)
 
+    def _to_lst(self, lst: Any) -> List[Any]:
+        # Convert to list, internal
+        if isinstance(lst, (list, tuple, set)):
+            pass
+        elif isinstance(lst, PdfObject):
+            lst = [lst]
+        elif isinstance(lst, str):
+            if lst[0] == "/":
+                lst = [NameObject(lst)]
+            else:
+                lst = [TextStringObject(lst)]
+        elif isinstance(lst, bytes):
+            lst = [ByteStringObject(lst)]
+        else:  # for numbers,...
+            lst = [lst]
+        return lst
+
+    def __add__(self, lst: Any) -> "ArrayObject":
+        """
+        Allow extension by adding list or add one element only
+
+        Args:
+            lst: any list, tuples are extended the list.
+            other types(numbers,...) will be appended.
+            if str is passed it will be converted into TextStringObject
+            or NameObject (if starting with "/")
+            if bytes is passed it will be converted into ByteStringObject
+
+        Returns:
+            ArrayObject with all elements
+        """
+        temp = ArrayObject(self)
+        temp.extend(self._to_lst(lst))
+        return temp
+
+    def __iadd__(self, lst: Any) -> Self:
+        """
+         Allow extension by adding list or add one element only
+
+        Args:
+            lst: any list, tuples are extended the list.
+            other types(numbers,...) will be appended.
+            if str is passed it will be converted into TextStringObject
+            or NameObject (if starting with "/")
+            if bytes is passed it will be converted into ByteStringObject
+        """
+        self.extend(self._to_lst(lst))
+        return self
+
+    def __isub__(self, lst: Any) -> Self:
+        """Allow to remove items"""
+        for x in self._to_lst(lst):
+            try:
+                x = self.index(x)
+                del self[x]
+            except ValueError:
+                pass
+        return self
+
     def write_to_stream(
         self, stream: StreamType, encryption_key: Union[None, str, bytes] = None
     ) -> None:

pyproject.toml

@@ -29,7 +29,7 @@ classifiers = [
 ]
 
 dependencies = [
-    "typing_extensions >= 3.7.4.3; python_version < '3.10'",
+    "typing_extensions >= 4.0; python_version < '3.11'",
     "dataclasses; python_version < '3.7'",
 ]
 

tests/test_generic.py

@@ -1278,3 +1278,32 @@ def test_indirect_object_page_dimensions():
     reader = PdfReader(data, strict=False)
     mediabox = reader.pages[0].mediabox
     assert mediabox == RectangleObject((0, 0, 792, 612))
+
+
+def test_array_operators():
+    a = ArrayObject(
+        [
+            NumberObject(1),
+            NumberObject(2),
+            NumberObject(3),
+            NumberObject(4),
+        ]
+    )
+    b = a + 5
+    assert isinstance(b, ArrayObject)
+    assert b == [1, 2, 3, 4, 5]
+    assert a == [1, 2, 3, 4]
+    a -= 2
+    a += "abc"
+    a -= (3, 4)
+    a += ["d", "e"]
+    a += BooleanObject(True)
+    assert a == [1, "abc", "d", "e", True]
+    a += "/toto"
+    assert isinstance(a[-1], NameObject)
+    assert isinstance(a[1], TextStringObject)
+    a += b"1234"
+    assert a[-1] == ByteStringObject(b"1234")
+    la = len(a)
+    a -= 300
+    assert len(a) == la