Add a lot of type annotations

src/ocrmypdf/_graft.py

@@ -18,6 +18,7 @@
 import logging
 from contextlib import suppress
 from pathlib import Path
+from typing import Optional
 
 import pikepdf
 
@@ -108,7 +109,12 @@ class OcrGrafter:
         self.interim_count = 0
 
     def graft_page(
-        self, *, pageno: int, image: Path, textpdf: Path, autorotate_correction: int
+        self,
+        *,
+        pageno: int,
+        image: Optional[Path],
+        textpdf: Optional[Path],
+        autorotate_correction: int,
     ):
         if textpdf and not self.font:
             self.font, self.font_key = self._find_font(textpdf)

src/ocrmypdf/_pipeline.py

@@ -20,7 +20,9 @@ import os
 import re
 import sys
 from datetime import datetime, timezone
+from pathlib import Path
 from shutil import copyfileobj
+from typing import Dict, Iterable, Optional
 
 import img2pdf
 import pikepdf
@@ -29,6 +31,7 @@ from PIL import Image, ImageColor, ImageDraw
 
 from ocrmypdf import leptonica
 from ocrmypdf._exec import unpaper
+from ocrmypdf._jobcontext import PageContext, PdfContext
 from ocrmypdf._version import PROGRAM_NAME
 from ocrmypdf._version import __version__ as VERSION
 from ocrmypdf.exceptions import (
@@ -170,7 +173,7 @@ def get_pdfinfo(
         raise InputFileError()
 
 
-def validate_pdfinfo_options(context):
+def validate_pdfinfo_options(context: PdfContext):
     pdfinfo = context.pdfinfo
     options = context.options
 
@@ -255,7 +258,7 @@ def get_canvas_square_dpi(pageinfo, options) -> Resolution:
     return Resolution(units, units)
 
 
-def is_ocr_required(page_context):
+def is_ocr_required(page_context: PageContext):
     pageinfo = page_context.pageinfo
     options = page_context.options
 
@@ -329,7 +332,7 @@ def is_ocr_required(page_context):
     return ocr_required
 
 
-def rasterize_preview(input_file, page_context):
+def rasterize_preview(input_file: Path, page_context: PageContext):
     output_file = page_context.get_path('rasterize_preview.jpg')
     canvas_dpi = get_canvas_square_dpi(page_context.pageinfo, page_context.options)
     page_dpi = get_page_square_dpi(page_context.pageinfo, page_context.options)
@@ -344,7 +347,7 @@ def rasterize_preview(input_file, page_context):
     return output_file
 
 
-def describe_rotation(page_context, orient_conf, correction):
+def describe_rotation(page_context: PageContext, orient_conf, correction: int):
     """
     Describe the page rotation we are going to perform.
     """
@@ -373,7 +376,7 @@ def describe_rotation(page_context, orient_conf, correction):
     return f"{facing}, confidence {orient_conf.confidence:.2f} - {action}"
 
 
-def get_orientation_correction(preview, page_context):
+def get_orientation_correction(preview: Path, page_context: PageContext):
     """Work out orientation correct for each page.
 
     We ask Ghostscript to draw a preview page, which will rasterize with the
@@ -405,7 +408,11 @@ def get_orientation_correction(preview, page_context):
 
 
 def rasterize(
-    input_file, page_context, correction=0, output_tag='', remove_vectors=None
+    input_file: Path,
+    page_context: PageContext,
+    correction: int = 0,
+    output_tag: str = '',
+    remove_vectors=None,
 ):
     colorspaces = ['pngmono', 'pnggray', 'png256', 'png16m']
     device_idx = 0
@@ -455,7 +462,7 @@ def rasterize(
     return output_file
 
 
-def preprocess_remove_background(input_file, page_context):
+def preprocess_remove_background(input_file: Path, page_context: PageContext):
     if any(image.bpc > 1 for image in page_context.pageinfo.images):
         output_file = page_context.get_path('pp_rm_bg.png')
         leptonica.remove_background(input_file, output_file)
@@ -465,21 +472,21 @@ def preprocess_remove_background(input_file, page_context):
         return input_file
 
 
-def preprocess_deskew(input_file, page_context):
+def preprocess_deskew(input_file: Path, page_context: PageContext):
     output_file = page_context.get_path('pp_deskew.png')
     dpi = get_page_square_dpi(page_context.pageinfo, page_context.options)
     leptonica.deskew(input_file, output_file, dpi.x)
     return output_file
 
 
-def preprocess_clean(input_file, page_context):
+def preprocess_clean(input_file: Path, page_context: PageContext):
     output_file = page_context.get_path('pp_clean.png')
     dpi = get_page_square_dpi(page_context.pageinfo, page_context.options)
     unpaper.clean(input_file, output_file, dpi.x, page_context.options.unpaper_args)
     return output_file
 
 
-def create_ocr_image(image, page_context):
+def create_ocr_image(image: Path, page_context: PageContext):
     """Create the image we send for OCR. May not be the same as the display
     image depending on preprocessing. This image will never be shown to the
     user."""
@@ -538,7 +545,7 @@ def create_ocr_image(image, page_context):
     return output_file
 
 
-def ocr_engine_hocr(input_file, page_context):
+def ocr_engine_hocr(input_file: Path, page_context: PageContext):
     hocr_out = page_context.get_path('ocr_hocr.hocr')
     hocr_text_out = page_context.get_path('ocr_hocr.txt')
     options = page_context.options
@@ -558,7 +565,7 @@ def should_visible_page_image_use_jpg(pageinfo):
     return pageinfo.images and all(im.enc == Encoding.jpeg for im in pageinfo.images)
 
 
-def create_visible_page_jpg(image, page_context):
+def create_visible_page_jpg(image: Path, page_context: PageContext) -> Path:
     output_file = page_context.get_path('visible.jpg')
     with Image.open(image) as im:
         # At this point the image should be a .png, but deskew, unpaper
@@ -577,7 +584,7 @@ def create_visible_page_jpg(image, page_context):
     return output_file
 
 
-def create_pdf_page_from_image(image, page_context):
+def create_pdf_page_from_image(image: Path, page_context: PageContext):
     # We rasterize a square DPI version of each page because most image
     # processing tools don't support rectangular DPI. Use the square DPI as it
     # accurately describes the image. It would be possible to resample the image
@@ -598,7 +605,7 @@ def create_pdf_page_from_image(image, page_context):
     return output_file
 
 
-def render_hocr_page(hocr, page_context):
+def render_hocr_page(hocr: Path, page_context: PageContext):
     output_file = page_context.get_path('ocr_hocr.pdf')
     dpi = get_page_square_dpi(page_context.pageinfo, page_context.options)
     hocrtransform = HocrTransform(hocr, dpi.x)  # square
@@ -612,7 +619,7 @@ def render_hocr_page(hocr, page_context):
     return output_file
 
 
-def ocr_engine_textonly_pdf(input_image, page_context):
+def ocr_engine_textonly_pdf(input_image: Path, page_context: PageContext):
     output_pdf = page_context.get_path('ocr_tess.pdf')
     output_text = page_context.get_path('ocr_tess.txt')
     options = page_context.options
@@ -627,7 +634,7 @@ def ocr_engine_textonly_pdf(input_image, page_context):
     return (output_pdf, output_text)
 
 
-def get_docinfo(base_pdf, context):
+def get_docinfo(base_pdf: pikepdf.Pdf, context: PdfContext) -> Dict[str, str]:
     options = context.options
 
     def from_document_info(key):
@@ -664,13 +671,13 @@ def get_docinfo(base_pdf, context):
     return pdfmark
 
 
-def generate_postscript_stub(context):
+def generate_postscript_stub(context: PdfContext):
     output_file = context.get_path('pdfa.ps')
     generate_pdfa_ps(output_file)
     return output_file
 
 
-def convert_to_pdfa(input_pdf, input_ps_stub, context):
+def convert_to_pdfa(input_pdf: Path, input_ps_stub: Path, context: PdfContext):
     options = context.options
     input_pdfinfo = context.pdfinfo
     fix_docinfo_file = context.get_path('fix_docinfo.pdf')
@@ -712,14 +719,14 @@ def convert_to_pdfa(input_pdf, input_ps_stub, context):
     return output_file
 
 
-def should_linearize(working_file, context):
+def should_linearize(working_file: Path, context: PdfContext):
     filesize = os.stat(working_file).st_size
     if filesize > (context.options.fast_web_view * 1_000_000):
         return True
     return False
 
 
-def metadata_fixup(working_file, context):
+def metadata_fixup(working_file: Path, context: PdfContext):
     output_file = context.get_path('metafix.pdf')
     options = context.options
 
@@ -768,7 +775,7 @@ def metadata_fixup(working_file, context):
     return output_file
 
 
-def optimize_pdf(input_file, context):
+def optimize_pdf(input_file: Path, context: PdfContext):
     output_file = context.get_path('optimize.pdf')
     save_settings = dict(
         compress_streams=True,
@@ -780,7 +787,7 @@ def optimize_pdf(input_file, context):
     return output_file
 
 
-def merge_sidecars(txt_files, context):
+def merge_sidecars(txt_files: Iterable[Optional[Path]], context: PdfContext):
     output_file = context.get_path('sidecar.txt')
     with open(output_file, 'w', encoding="utf-8") as stream:
         for page_num, txt_file in enumerate(txt_files):
@@ -801,7 +808,7 @@ def merge_sidecars(txt_files, context):
     return output_file
 
 
-def copy_final(input_file, output_file, _context):
+def copy_final(input_file: Path, output_file: Path, _context: PdfContext):
     log.debug('%s -> %s', input_file, output_file)
     with open(input_file, 'rb') as input_stream:
         if output_file == '-':

src/ocrmypdf/_plugin_manager.py

@@ -20,7 +20,7 @@ import importlib
 import importlib.util
 import sys
 from pathlib import Path
-from typing import List
+from typing import List, Tuple
 
 import pluggy
 
@@ -56,7 +56,7 @@ def get_plugin_manager(plugins: List[str], builtins=True):
 
 def get_parser_options_plugins(
     args,
-) -> (argparse.ArgumentParser, argparse.Namespace, pluggy.PluginManager):
+) -> Tuple[argparse.ArgumentParser, argparse.Namespace, pluggy.PluginManager]:
     pre_options, _unused = plugins_only_parser.parse_known_args(args=args)
     plugin_manager = get_plugin_manager(pre_options.plugins)
 

src/ocrmypdf/_sync.py

@@ -20,7 +20,6 @@ import logging.handlers
 import os
 import sys
 import threading
-from collections import namedtuple
 from functools import partial
 from pathlib import Path
 from tempfile import mkdtemp
@@ -68,9 +67,14 @@ from ocrmypdf.pdfa import file_claims_pdfa
 
 log = logging.getLogger(__name__)
 
-PageResult = namedtuple(
-    'PageResult', 'pageno, pdf_page_from_image, ocr, text, orientation_correction'
-)
+
+class PageResult(NamedTuple):
+    pageno: int
+    pdf_page_from_image: Optional[Path]
+    ocr: Optional[Path]
+    text: Optional[Path]
+    orientation_correction: int
+
 
 tls = threading.local()
 tls.pageno = None
@@ -221,7 +225,7 @@ def exec_page_sync(page_context: PageContext):
     )
 
 
-def post_process(pdf_file, context):
+def post_process(pdf_file, context: PdfContext):
     pdf_out = pdf_file
     if context.options.output_type.startswith('pdfa'):
         ps_stub_out = generate_postscript_stub(context)
@@ -231,14 +235,14 @@ def post_process(pdf_file, context):
     return optimize_pdf(pdf_out, context)
 
 
-def worker_init(max_pixels):
+def worker_init(max_pixels: int):
     # In Windows, child process will not inherit our change to this value in
     # the parent process, so ensure workers get it set. Not needed when running
     # threaded, but harmless to set again.
     PIL.Image.MAX_IMAGE_PIXELS = max_pixels
 
 
-def exec_concurrent(context):
+def exec_concurrent(context: PdfContext):
     """Execute the pipeline concurrently"""
 
     # Run exec_page_sync on every page context
@@ -246,10 +250,10 @@ def exec_concurrent(context):
     if max_workers > 1:
         log.info("Start processing %d pages concurrently", max_workers)
 
-    sidecars = [None] * len(context.pdfinfo)
+    sidecars: List[Optional[Path]] = [None] * len(context.pdfinfo)
     ocrgraft = OcrGrafter(context)
 
-    def update_page(result, pbar):
+    def update_page(result: PageResult, pbar):
         sidecars[result.pageno] = result.text
         pbar.update()
         ocrgraft.graft_page(

src/ocrmypdf/_validation.py

@@ -24,6 +24,7 @@ import sys
 import unicodedata
 from pathlib import Path
 from shutil import copyfileobj
+from typing import Tuple
 
 import pikepdf
 import PIL
@@ -329,7 +330,7 @@ def log_page_orientations(pdfinfo):
         log.info('Page orientations detected: %s', ' '.join(orientations))
 
 
-def create_input_file(options, work_folder: Path) -> (Path, str):
+def create_input_file(options, work_folder: Path) -> Tuple[Path, str]:
     if options.input_file == '-':
         # stdin
         log.info('reading file from standard input')

src/ocrmypdf/api.py

@@ -264,6 +264,8 @@ def ocr(  # pylint: disable=unused-argument
     """
     if not plugins:
         plugins = []
+    else:
+        plugins = list(plugins)
 
     parser = get_parser()
     _plugin_manager = get_plugin_manager(plugins)

src/ocrmypdf/helpers.py

@@ -27,6 +27,7 @@ from functools import wraps
 from io import StringIO
 from math import isclose
 from pathlib import Path
+from typing import Any, Sequence
 
 import pikepdf
 
@@ -104,28 +105,28 @@ def safe_symlink(input_file: os.PathLike, soft_link_name: os.PathLike):
     os.symlink(os.path.abspath(input_file), soft_link_name)
 
 
-def samefile(f1, f2):
+def samefile(f1: os.PathLike, f2: os.PathLike):
     if os.name == 'nt':
         return f1 == f2
     else:
         return os.path.samefile(f1, f2)
 
 
-def is_iterable_notstr(thing):
+def is_iterable_notstr(thing: Any) -> bool:
     return isinstance(thing, Iterable) and not isinstance(thing, str)
 
 
-def monotonic(L: Iterable):
+def monotonic(L: Sequence) -> bool:
     """Does list increase monotonically?"""
     return all(b > a for a, b in zip(L, L[1:]))
 
 
-def page_number(input_file: os.PathLike):
+def page_number(input_file: os.PathLike) -> int:
     """Get one-based page number implied by filename (000002.pdf -> 2)"""
     return int(os.path.basename(os.fspath(input_file))[0:6])
 
 
-def available_cpu_count():
+def available_cpu_count() -> int:
     try:
         return multiprocessing.cpu_count()
     except NotImplementedError:
@@ -136,7 +137,7 @@ def available_cpu_count():
     return 1
 
 
-def is_file_writable(test_file: os.PathLike):
+def is_file_writable(test_file: os.PathLike) -> bool:
     """Intentionally racy test if target is writable.
 
     We intend to write to the output file if and only if we succeed and
@@ -171,7 +172,7 @@ def is_file_writable(test_file: os.PathLike):
         return False
 
 
-def check_pdf(input_file):
+def check_pdf(input_file: Path) -> bool:
     pdf = None
     try:
         pdf = pikepdf.open(input_file)

src/ocrmypdf/hocrtransform.py

@@ -34,6 +34,7 @@ import re
 from collections import namedtuple
 from math import atan, cos, sin
 from pathlib import Path
+from typing import Union
 from xml.etree import ElementTree
 
 from reportlab.lib.units import inch
@@ -66,9 +67,9 @@ class HocrTransform:
         {'ff': 'ff', 'ffi': 'f‌f‌i', 'ffl': 'f‌f‌l', 'fi': 'fi', 'fl': 'fl'}
     )
 
-    def __init__(self, hocr_filename: str, dpi: float):
+    def __init__(self, hocr_filename: Union[str, Path], dpi: float):
         self.dpi = dpi
-        self.hocr = ElementTree.parse(hocr_filename)
+        self.hocr = ElementTree.parse(os.fspath(hocr_filename))
 
         # if the hOCR file has a namespace, ElementTree requires its use to
         # find elements

src/ocrmypdf/pluginspec.py

@@ -19,13 +19,17 @@ from abc import ABC, abstractmethod, abstractstaticmethod
 from argparse import ArgumentParser, Namespace
 from collections import namedtuple
 from pathlib import Path
-from typing import AbstractSet, List, Optional
+from typing import TYPE_CHECKING, AbstractSet, List, Optional
 
 import pluggy
 from PIL import Image
 
 from ocrmypdf.helpers import Resolution
 
+if TYPE_CHECKING:
+    from ocrmypdf._jobcontext import PageContext
+    from ocrmypdf.pdfinfo import PdfInfo
+
 hookspec = pluggy.HookspecMarker('ocrmypdf')
 
 # pylint: disable=unused-argument