From b4b32a35b5d6957a534d05999ed454c9a588093f Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <james@purplerock.ca>
Date: Wed, 15 Sep 2021 00:09:08 -0700
Subject: [PATCH] optimize: fix typing consistency

---
 src/ocrmypdf/optimize.py | 36 ++++++++++++++++++++++--------------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/src/ocrmypdf/optimize.py b/src/ocrmypdf/optimize.py
index 0df62327..9712983a 100644
--- a/src/ocrmypdf/optimize.py
+++ b/src/ocrmypdf/optimize.py
@@ -25,8 +25,16 @@ from typing import (
 )
 
 import img2pdf
-import pikepdf
-from pikepdf import Dictionary, Name, Object, Pdf, PdfImage
+from pikepdf import (
+    Dictionary,
+    Name,
+    Object,
+    ObjectStreamMode,
+    Pdf,
+    PdfImage,
+    Stream,
+    UnsupportedImageTypeError,
+)
 from PIL import Image
 
 from ocrmypdf import leptonica
@@ -63,7 +71,7 @@ def jpg_name(root: Path, xref: Xref) -> Path:
 
 
 def extract_image_filter(
-    pike: Pdf, root: Path, image: Object, xref: Xref
+    pike: Pdf, root: Path, image: Stream, xref: Xref
 ) -> Optional[Tuple[PdfImage, Tuple[Name, Object]]]:
     del pike  # unused args
     del root
@@ -104,7 +112,7 @@ def extract_image_filter(
 
 
 def extract_image_jbig2(
-    *, pike: pikepdf.Pdf, root: Path, image: Object, xref: Xref, options
+    *, pike: Pdf, root: Path, image: Stream, xref: Xref, options
 ) -> Optional[XrefExt]:
     del options  # unused arg
 
@@ -123,16 +131,16 @@ def extract_image_jbig2(
         # Showing the palette or ICC to jbig2enc will cause it to perform
         # colorspace transform to 1bpp, which will conflict the palette or
         # ICC if it exists.
-        colorspace = pim.obj.get(pikepdf.Name.ColorSpace, None)
+        colorspace = pim.obj.get(Name.ColorSpace, None)
         if colorspace is not None or pim.image_mask:
             try:
                 # Set to DeviceGray temporarily; we already in 1 bpc.
-                pim.obj.ColorSpace = pikepdf.Name.DeviceGray
+                pim.obj.ColorSpace = Name.DeviceGray
                 imgname = root / f'{xref:08d}'
                 with imgname.open('wb') as f:
                     ext = pim.extract_to(stream=f)
                 imgname.rename(imgname.with_suffix(ext))
-            except pikepdf.UnsupportedImageTypeError:
+            except UnsupportedImageTypeError:
                 return None
             finally:
                 # Restore image colorspace after temporarily setting it to DeviceGray
@@ -145,7 +153,7 @@ def extract_image_jbig2(
 
 
 def extract_image_generic(
-    *, pike: Pdf, root: Path, image: PdfImage, xref: Xref, options
+    *, pike: Pdf, root: Path, image: Stream, xref: Xref, options
 ) -> Optional[XrefExt]:
     result = extract_image_filter(pike, root, image, xref)
     if result is None:
@@ -178,7 +186,7 @@ def extract_image_generic(
             with imgname.open('wb') as f:
                 ext = pim.extract_to(stream=f)
             imgname.rename(imgname.with_suffix(ext))
-        except pikepdf.UnsupportedImageTypeError:
+        except UnsupportedImageTypeError:
             return None
         return XrefExt(xref, ext)
     elif (
@@ -374,7 +382,7 @@ def convert_to_jbig2(
         jbig2_symfile = root / (prefix + '.sym')
         if jbig2_symfile.exists():
             jbig2_globals_data = jbig2_symfile.read_bytes()
-            jbig2_globals = pikepdf.Stream(pike, jbig2_globals_data)
+            jbig2_globals = Stream(pike, jbig2_globals_data)
             jbig2_globals_dict = Dictionary(JBIG2Globals=jbig2_globals)
         elif options.jbig2_page_group_size == 1:
             jbig2_globals_dict = None
@@ -445,7 +453,7 @@ def _transcode_png(pike: Pdf, filename: Path, xref: Xref) -> bool:
     with output.open('wb') as f:
         img2pdf.convert(fspath(filename), outputstream=f)
 
-    with pikepdf.open(output) as pdf_image:
+    with Pdf.open(output) as pdf_image:
         foreign_image = next(pdf_image.pages[0].images.values())
         local_image = pike.copy_foreign(foreign_image)
 
@@ -544,7 +552,7 @@ def optimize(
     if options.jbig2_page_group_size == 0:
         options.jbig2_page_group_size = 10 if options.jbig2_lossy else 1
 
-    with pikepdf.Pdf.open(input_file) as pike:
+    with Pdf.open(input_file) as pike:
         root = output_file.parent / 'images'
         root.mkdir(exist_ok=True)
 
@@ -576,7 +584,7 @@ def optimize(
     if savings < 0:
         log.info("Image optimization did not improve the file - discarded")
         # We still need to save the file
-        with pikepdf.open(input_file) as pike:
+        with Pdf.open(input_file) as pike:
             pike.remove_unreferenced_resources()
             pike.save(output_file, **save_settings)
     else:
@@ -623,7 +631,7 @@ def main(infile, outfile, level, jobs=1):
             dict(
                 compress_streams=True,
                 preserve_pdfa=True,
-                object_stream_mode=pikepdf.ObjectStreamMode.generate,
+                object_stream_mode=ObjectStreamMode.generate,
             ),
         )
         copy(fspath(tmpout), fspath(outfile))