From 8d12ecb798af2c5681b0b7114fb5505539317130 Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <james@purplerock.ca>
Date: Sun, 24 Sep 2023 20:21:26 -0700
Subject: [PATCH] Refactor docinfo repair code

---
 src/ocrmypdf/_pipeline.py | 34 ++++++++++++++++++++--------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/src/ocrmypdf/_pipeline.py b/src/ocrmypdf/_pipeline.py
index 01d23003..fa7288db 100644
--- a/src/ocrmypdf/_pipeline.py
+++ b/src/ocrmypdf/_pipeline.py
@@ -787,21 +787,8 @@ def convert_to_pdfa(input_pdf: Path, input_ps_stub: Path, context: PdfContext) -
     # NULs in DocumentInfo seem to be common since older Acrobats included them.
     # pikepdf can deal with this, but we make the world a better place by
     # stamping them out as soon as possible.
-    modified = False
     with pikepdf.open(input_pdf) as pdf_file:
-        try:
-            len(pdf_file.docinfo)
-        except TypeError:
-            log.error(
-                "File contains a malformed DocumentInfo block - continuing anyway"
-            )
-        else:
-            if pdf_file.docinfo:
-                for k, v in pdf_file.docinfo.items():
-                    if b'\x00' in bytes(v):
-                        pdf_file.docinfo[k] = bytes(v).replace(b'\x00', b'')
-                        modified = True
-        if modified:
+        if _repair_docinfo_nulls(pdf_file):
             pdf_file.save(fix_docinfo_file)
         else:
             safe_symlink(input_pdf, fix_docinfo_file)
@@ -824,6 +811,25 @@ def convert_to_pdfa(input_pdf: Path, input_ps_stub: Path, context: PdfContext) -
     return output_file
 
 
+def _repair_docinfo_nulls(pdf):
+    """If the DocumentInfo block contains NUL characters, remove them.
+
+    If the DocumentInfo block is malformed, log an error and continue.
+    """
+    modified = False
+    try:
+        if not isinstance(pdf.docinfo, pikepdf.Dictionary):
+            raise TypeError("DocumentInfo is not a dictionary")
+        for k, v in pdf.docinfo.items():
+            if isinstance(v, str) and b'\x00' in bytes(v):
+                pdf.docinfo[k] = bytes(v).replace(b'\x00', b'')
+                modified = True
+    except TypeError:
+        # TypeError can also be raised if dictionary items are unexpected types
+        log.error("File contains a malformed DocumentInfo block - continuing anyway.")
+    return modified
+
+
 def should_linearize(working_file: Path, context: PdfContext) -> bool:
     filesize = os.stat(working_file).st_size
     if filesize > (context.options.fast_web_view * 1_000_000):