ENH: Automatically preserve links in added pages

Author: larsga

pypdf/_writer.py

@@ -146,6 +146,54 @@ def _rolling_checksum(stream: BytesIO, blocksize: int = 65536) -> str:
     return hash.hexdigest()
 
 
+class NamedRefLink:
+    "Named reference link being preserved until we can resolve it correctly."
+
+    def __init__(self, ref: TextStringObject, source_pdf: Union[PdfReader,None]):
+        "ref: TextStringObject with named reference"
+        self._ref = ref
+        self._source_pdf = source_pdf
+
+    def find_referenced_page(self) -> Union[IndirectObject,None]:
+        if not self._source_pdf:
+            return None
+        dest = self._source_pdf.named_destinations.get(str(self._ref))
+        if dest:
+            return dest.page
+
+    def patch_reference(self, target_pdf, new_page: IndirectObject) -> None:
+        "target_pdf: PdfWriter which the new link went into"
+
+        def _get_referenced_page_ix(pdf: PdfWriter, ref: IndirectObject):
+            try:
+                return [
+                    p.indirect_reference for p in pdf.flattened_pages
+                ].index(ref)
+            except ValueError:
+                return None
+
+        # point named destination in new PDF to the new page
+        if str(self._ref) not in target_pdf.named_destinations:
+            new_page_ix = _get_referenced_page_ix(target_pdf, new_page)
+            if new_page_ix is not None:
+                target_pdf.add_named_destination(str(self._ref), new_page_ix)
+
+
+class DirectRefLink:
+    "Direct reference link being preserved until we can resolve it correctly."
+
+    def __init__(self, ref: ArrayObject, source_pdf: Union[PdfReader,None]):
+        "ref: an ArrayObject whose first element is the Page indir obj"
+        self._ref = ref
+
+    def find_referenced_page(self) -> IndirectObject:
+        return self._ref[0]
+
+    def patch_reference(self, target_pdf, new_page: IndirectObject) -> None:
+        "target_pdf: PdfWriter which the new link went into"
+        self._ref[0] = new_page
+
+
 class PdfWriter(PdfDocCommon):
     """
     Write a PDF file out, given pages produced by another class or through
@@ -209,6 +257,11 @@ def __init__(
         """The PDF file identifier,
         defined by the ID in the PDF file's trailer dictionary."""
 
+        self._unresolved_links: list[Union[NamedRefLink,DirectRefLink]] = []
+        "Tracks links in pages added to the writer for resolving later."
+        self._merged_in_pages: List[Tuple[IndirectObject,IndirectObject]] = []
+        "Tracks pages added to the writer and what page they turned into."
+
         if self.incremental:
             if isinstance(fileobj, (str, Path)):
                 with open(fileobj, "rb") as f:
@@ -482,12 +535,50 @@ def _add_page(
             ]
         except Exception:
             pass
+
+        def _extract_links(new_page: PageObject, old_page: PageObject) -> List[Union[NamedRefLink,DirectRefLink]]:
+            new_links = [_build_link(link, new_page) for link in new_page.get("/Annots", [])]
+            old_links = [_build_link(link, old_page) for link in old_page.get("/Annots", [])]
+
+            return [(new_link, old_link) for (new_link, old_link)
+                    in zip(new_links, old_links)
+                    if new_link]
+
+        def _build_link(indir_obj: IndirectObject, page: PageObject) -> Union[NamedRefLink,DirectRefLink]:
+            if page.pdf:
+                src = page.pdf
+            else:
+                src = None
+
+            link = indir_obj.get_object()
+            if link["/Subtype"] != "/Link":
+                return None
+
+            if "/A" in link:
+                action = link["/A"]
+                if action.get("/S") != "/GoTo":
+                    return None
+
+                return _create_link(action["/D"], src)
+
+            elif "/Dest" in link:
+                return _create_link(link["/Dest"], src)
+
+            return None # nothing we need to do
+
+        def _create_link(ref: Union[TextStringObject,ArrayObject], src: Union[PdfReader,None]) -> Union[NamedRefLink,DirectRefLink]:
+            if isinstance(ref, TextStringObject):
+                return NamedRefLink(ref, src)
+            else:
+                return DirectRefLink(ref, src)
+
         page = cast(
             "PageObject", page_org.clone(self, False, excluded_keys).get_object()
         )
         if page_org.pdf is not None:
             other = page_org.pdf.pdf_header
             self.pdf_header = _get_max_pdf_version_header(self.pdf_header, other)
+
         node, idx = self._get_page_in_node(index)
         page[NameObject(PA.PARENT)] = node.indirect_reference
 
@@ -505,6 +596,16 @@ def _add_page(
             recurse += 1
             if recurse > 1000:
                 raise PyPdfError("Too many recursive calls!")
+
+        if page_org.pdf is not None:
+            # the page may contain links to other pages, and those other
+            # pages may or may not already be added.  we store the
+            # information we need, so that we can resolve the references
+            # later.
+            self._unresolved_links.extend(_extract_links(page, page_org))
+            self._merged_in_pages.append( (page_org.indirect_reference,
+                                           page.indirect_reference) )
+
         return page
 
     def set_need_appearances_writer(self, state: bool = True) -> None:
@@ -1349,6 +1450,30 @@ def encrypt(
             self._add_object(entry)
         self._encrypt_entry = entry
 
+    def _resolve_links(self) -> None:
+        """Patch up links that were added to the document earlier, to
+        make sure they still point to the same pages."""
+
+        def _get_referenced_page_ix(pdf: PdfWriter, ref: IndirectObject) -> Union[int,None]:
+            try:
+                return [
+                    p.indirect_reference for p in pdf.flattened_pages
+                ].index(ref)
+            except ValueError:
+                return None
+
+        for (new_link, old_link) in self._unresolved_links:
+            old_page = old_link.find_referenced_page()
+            if not old_page:
+                continue
+            new_page = None
+            for (page_org, page_created) in self._merged_in_pages:
+                if page_org == old_page:
+                    new_page = page_created
+            if new_page is None:
+                continue
+            new_link.patch_reference(self, new_page)
+
     def write_stream(self, stream: StreamType) -> None:
         if hasattr(stream, "mode") and "b" not in stream.mode:
             logger_warning(
@@ -1360,6 +1485,7 @@ def write_stream(self, stream: StreamType) -> None:
         # if not self._root:
         #   self._root = self._add_object(self._root_object)
         # self._sweep_indirect_references(self._root)
+        self._resolve_links()
 
         if self.incremental:
             self._reader.stream.seek(0)

tests/example_files.yaml

@@ -110,3 +110,7 @@
   url: https://github.com/py-pdf/pypdf/files/12483807/AEO.1172.pdf
 - local_filename: iss3268.pdf
   url: https://github.com/user-attachments/files/20060394/broken.pdf
+- local_filename: direct-link.pdf
+  url: https://github.com/user-attachments/files/20348304/tst.pdf
+- local_filename: named-reference.pdf
+  url: https://github.com/user-attachments/files/20455804/MinimalJob.pdf

tests/test_merger.py

@@ -409,3 +409,92 @@ def test_deprecate_pdfmerger():
 def test_get_reference():
     writer = PdfWriter(RESOURCE_ROOT / "crazyones.pdf")
     assert writer.get_reference(writer.pages[0]) == writer.pages[0].indirect_reference
+
+
+@pytest.mark.enable_socket
+def test_direct_link_preserved(pdf_file_path):
+    # this could be any PDF -- we don't care which
+    reader = PdfReader(BytesIO(get_data_from_url(name="iss3268.pdf")))
+    writer = PdfWriter(clone_from = reader)
+
+    # this PDF has a direct link from p1 to p2
+    merger = PdfReader(BytesIO(get_data_from_url(name="direct-link.pdf")))
+    for p in merger.pages:
+        writer.add_page(p)
+
+    writer.write(pdf_file_path)
+
+    check = PdfReader(pdf_file_path)
+    page3 = check.pages[2]
+    link = page3["/Annots"][0].get_object()
+    assert "/Link" == link["/Subtype"]
+    dest = link["/Dest"][0] # indirect ref of page referred to
+
+    page4 = check.flattened_pages[3]
+    assert dest == page4.indirect_reference, "Link from page 3 to page 4 is broken"
+
+
+@pytest.mark.enable_socket
+def test_direct_link_preserved_reordering(pdf_file_path):
+    # this could be any PDF -- we don't care which
+    reader = PdfReader(BytesIO(get_data_from_url(name="iss3268.pdf")))
+    writer = PdfWriter(clone_from = reader)
+
+    # this PDF has a direct link from p1 to p2
+    merger = PdfReader(BytesIO(get_data_from_url(name="direct-link.pdf")))
+    for p in merger.pages:
+        writer.add_page(p)
+
+    # let's insert a page to mess up the page order
+    writer.insert_page(reader.pages[0], 3)
+
+    writer.write(pdf_file_path)
+
+    check = PdfReader(pdf_file_path)
+    page3 = check.pages[2]
+    link = page3["/Annots"][0].get_object()
+    assert "/Link" == link["/Subtype"]
+    dest = link["/Dest"][0] # indirect ref of page referred to
+
+    page5 = check.flattened_pages[4] # it moved one out
+    assert dest == page5.indirect_reference, "Link from page 3 to page 5 is broken"
+
+
+@pytest.mark.enable_socket
+def test_direct_link_page_missing(pdf_file_path):
+    # this could be any PDF -- we don't care which
+    reader = PdfReader(BytesIO(get_data_from_url(name="iss3268.pdf")))
+    writer = PdfWriter(clone_from = reader)
+
+    # this PDF has a direct link from p1 to p2
+    merger = PdfReader(BytesIO(get_data_from_url(name="direct-link.pdf")))
+    writer.add_page(merger.pages[0])
+    # but we're not adding page 2
+
+    writer.write(pdf_file_path) # verify nothing crashes
+
+
+@pytest.mark.enable_socket
+def test_named_reference_preserved(pdf_file_path):
+    # this could be any PDF -- we don't care which
+    reader = PdfReader(BytesIO(get_data_from_url(name="iss3268.pdf")))
+    writer = PdfWriter(clone_from = reader)
+
+    # this PDF has a named reference from from p3 to p5
+    merger = PdfReader(BytesIO(get_data_from_url(name="named-reference.pdf")))
+    for p in merger.pages:
+        writer.add_page(p)
+
+    writer.write(pdf_file_path)
+
+    check = PdfReader(pdf_file_path)
+    page5 = check.pages[4]
+    page7 = check.flattened_pages[6]
+    for link in page5["/Annots"]:
+        action = link["/A"]
+        assert "/GoTo" == action.get("/S")
+        dest = str(action["/D"])
+        assert dest in check.named_destinations
+        pref = check.named_destinations[dest].page
+
+        assert pref == page7.indirect_reference, "Link from page 5 to page 7 is broken"