ENH: Add page label support to PdfWriter

pypdf/_writer.py

@@ -123,6 +123,14 @@
     ZoomArgType,
 )
 
+try:
+    from typing import Literal  # type: ignore[attr-defined]
+except ImportError:
+    # PEP 586 introduced typing.Literal with Python 3.8
+    # For older Python versions, the backport typing_extensions is necessary:
+    from typing_extensions import Literal  # type: ignore[misc]
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -2874,6 +2882,113 @@ def reset_translation(
         else:
             raise Exception("invalid parameter {reader}")
 
+    def set_page_label(
+        self,
+        page_number_from: int,
+        page_number_to: int,
+        style: Optional[Literal["/D", "/R", "/r", "/A", "/a"]] = None,
+        prefix: Optional[str] = None,
+        start: Optional[int] = 0,
+    ):
+        if style is None and prefix is None:
+            raise ValueError("at least one between style and prefix must be given")
+        if page_number_from < 1:
+            raise ValueError("page_index_from must be equal or greater then 1")
+        if page_number_to < page_number_from:
+            raise ValueError(
+                "page_index_to must be equal or greater then page_index_from"
+            )
+        if page_number_to > len(self.pages):
+            raise ValueError("page_index_to exceeds number of pages")
+        if start != 0 and start < 1:
+            raise ValueError("start must be equal or greater than one")
+
+        self._set_page_label(
+            page_number_from - 1, page_number_to - 1, style, prefix, start
+        )
+
+    def _set_page_label(
+        self,
+        page_index_from: int,
+        page_index_to: int,
+        style: Optional[Literal["/D", "/R", "/r", "/A", "/a"]] = None,
+        prefix: Optional[str] = None,
+        start: Optional[int] = 0,
+    ):
+        default_page_label = DictionaryObject()
+        default_page_label[NameObject("/S")] = NameObject("/D")
+
+        new_page_label = DictionaryObject()
+        if style is not None:
+            new_page_label[NameObject("/S")] = NameObject(style)
+        if prefix is not None:
+            new_page_label[NameObject("/P")] = TextStringObject(prefix)
+        if start != 0:
+            new_page_label[NameObject("/St")] = NumberObject(start)
+
+        if not NameObject(CatalogDictionary.PAGE_LABELS) in self._root_object:
+            nums = ArrayObject()
+            self._nums_insert(NumberObject(0), default_page_label, nums)
+            page_labels = TreeObject()
+            page_labels[NameObject("/Nums")] = nums
+            self._root_object[NameObject(CatalogDictionary.PAGE_LABELS)] = page_labels
+
+        page_labels = self._root_object[NameObject(CatalogDictionary.PAGE_LABELS)]
+        nums = page_labels[NameObject("/Nums")]
+
+        self._nums_insert(NumberObject(page_index_from), new_page_label, nums)
+        self._nums_clear_range(NumberObject(page_index_from), page_index_to, nums)
+        next_label_pos, *_ = self._nums_next(NumberObject(page_index_from), nums)
+        if next_label_pos != page_index_to + 1 and page_index_to + 1 < len(self.pages):
+            self._nums_insert(NumberObject(page_index_to + 1), default_page_label, nums)
+
+        page_labels[NameObject("/Nums")] = nums
+        self._root_object[NameObject(CatalogDictionary.PAGE_LABELS)] = page_labels
+
+    def _nums_insert(
+        self,
+        key: NumberObject,
+        value: DictionaryObject,
+        nums: ArrayObject,
+    ):
+        if len(nums) % 2 != 0:
+            raise ValueError("nums must contain an even number of elements")
+
+        i = len(nums)
+        while i != 0 and key <= nums[i - 2]:
+            i = i - 2
+
+        if i < len(nums) and key == nums[i]:
+            nums[i + 1] = value
+        else:
+            nums.insert(i, key)
+            nums.insert(i + 1, value)
+
+    def _nums_clear_range(
+        self,
+        key: NumberObject,
+        page_index_to: int,
+        nums: ArrayObject,
+    ):
+        if page_index_to < key:
+            raise ValueError("page_index_to must be greater or equal than key")
+
+        i = nums.index(key) + 2
+        while i < len(nums) and nums[i] <= page_index_to:
+            nums.pop(i)
+            nums.pop(i)
+
+    def _nums_next(
+        self,
+        key: NumberObject,
+        nums: ArrayObject,
+    ):
+        i = nums.index(key) + 2
+        if i < len(nums):
+            return (nums[i], nums[i + 1])
+        else:
+            return (None, None)
+
 
 def _pdf_objectify(obj: Union[Dict[str, Any], str, int, List[Any]]) -> PdfObject:
     if isinstance(obj, PdfObject):

tests/test_writer.py

@@ -1008,3 +1008,98 @@ def test_append_multiple():
     pages = writer._root_object["/Pages"]["/Kids"]
     assert pages[0] not in pages[1:]  # page not repeated
     assert pages[-1] not in pages[0:-1]  # page not repeated
+
+
+@pytest.mark.samples
+def test_set_page_labels():
+    src = RESOURCE_ROOT / "GeoBase_NHNC1_Data_Model_UML_EN.pdf"  # File without labels
+    target = "pypdf-output.pdf"
+    reader = PdfReader(src)
+
+    expected = [
+        "i",
+        "ii",
+        "1",
+        "2",
+        "A",
+        "B",
+        "1",
+        "2",
+        "3",
+        "4",
+        "A",
+        "i",
+        "I",
+        "II",
+        "1",
+        "2",
+        "3",
+        "I",
+        "II",
+    ]
+
+    # Tests full lenght with labels assigned at first and last elements
+    # Tests different labels assigned to consecutive ranges
+    writer = PdfWriter()
+    writer.clone_document_from_reader(reader)
+    writer.set_page_label(1, 2, "/r")
+    writer.set_page_label(5, 6, "/A")
+    writer.set_page_label(11, 11, "/A")
+    writer.set_page_label(12, 12, "/r")
+    writer.set_page_label(13, 14, "/R")
+    writer.set_page_label(18, 19, "/R")
+    writer.write(target)
+    assert PdfReader(target).page_labels == expected
+
+    writer = PdfWriter()  # Same labels, different set order
+    writer.clone_document_from_reader(reader)
+    writer.set_page_label(18, 19, "/R")
+    writer.set_page_label(5, 6, "/A")
+    writer.set_page_label(11, 11, "/A")
+    writer.set_page_label(1, 2, "/r")
+    writer.set_page_label(13, 14, "/R")
+    writer.set_page_label(12, 12, "/r")
+    writer.write(target)
+    assert PdfReader(target).page_labels == expected
+
+    # Tests labels assigned only in the middle
+    # Tests label assigned to a range already containing labled ranges
+    expected = ["1", "2", "i", "ii", "iii", "iv", "v", "1"]
+    writer = PdfWriter()
+    writer.clone_document_from_reader(reader)
+    writer.set_page_label(4, 5, "/a")
+    writer.set_page_label(6, 6, "/A")
+    writer.set_page_label(3, 7, "/r")
+    writer.write(target)
+    assert PdfReader(target).page_labels[: len(expected)] == expected
+
+    # Tests labels assigned inside a previously existing range
+    expected = ["1", "2", "i", "a", "b", "A", "1", "1", "2"]
+    # Ones repeat because user didnt cover the entire original range
+    writer = PdfWriter()
+    writer.clone_document_from_reader(reader)
+    writer.set_page_label(3, 7, "/r")
+    writer.set_page_label(4, 5, "/a")
+    writer.set_page_label(6, 6, "/A")
+    writer.write(target)
+    assert PdfReader(target).page_labels[: len(expected)] == expected
+
+    src = SAMPLE_ROOT / "009-pdflatex-geotopo/GeoTopo.pdf"  # File with pre existing labels
+    target = "pypdf-output.pdf"
+    reader = PdfReader(src)
+
+    # Tests adding labels to existing ones
+    expected = ["i", "ii", "A", "B", "1"]
+    writer = PdfWriter()
+    writer.clone_document_from_reader(reader)
+    writer.set_page_label(3, 4, "/A")
+    writer.write(target)
+    assert PdfReader(target).page_labels[: len(expected)] == expected
+
+    # Tests replacing existing lables
+    expected = ["A", "B", "1", "1", "2"]
+    writer = PdfWriter()
+    writer.clone_document_from_reader(reader)
+    writer.set_page_label(1, 2, "/A")
+    writer.write(target)
+    assert PdfReader(target).page_labels[: len(expected)] == expected