From 698aab4f75b44836b778f8a2975bd9fe0ecae2c0 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Mon, 15 Jun 2020 15:10:59 -0700 Subject: [PATCH] Add a lot of type annotations --- src/ocrmypdf/_graft.py | 8 ++++- src/ocrmypdf/_pipeline.py | 53 +++++++++++++++++++-------------- src/ocrmypdf/_plugin_manager.py | 4 +-- src/ocrmypdf/_sync.py | 22 ++++++++------ src/ocrmypdf/_validation.py | 3 +- src/ocrmypdf/api.py | 2 ++ src/ocrmypdf/helpers.py | 15 +++++----- src/ocrmypdf/hocrtransform.py | 5 ++-- src/ocrmypdf/pluginspec.py | 6 +++- 9 files changed, 72 insertions(+), 46 deletions(-) diff --git a/src/ocrmypdf/_graft.py b/src/ocrmypdf/_graft.py index de9ebda0..14c09af1 100644 --- a/src/ocrmypdf/_graft.py +++ b/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) diff --git a/src/ocrmypdf/_pipeline.py b/src/ocrmypdf/_pipeline.py index 8c27127f..40a27e3c 100644 --- a/src/ocrmypdf/_pipeline.py +++ b/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 == '-': diff --git a/src/ocrmypdf/_plugin_manager.py b/src/ocrmypdf/_plugin_manager.py index baa25e4a..5bc1ab60 100644 --- a/src/ocrmypdf/_plugin_manager.py +++ b/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) diff --git a/src/ocrmypdf/_sync.py b/src/ocrmypdf/_sync.py index cdfff74e..b55ace4f 100644 --- a/src/ocrmypdf/_sync.py +++ b/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( diff --git a/src/ocrmypdf/_validation.py b/src/ocrmypdf/_validation.py index 06a16670..254410c1 100644 --- a/src/ocrmypdf/_validation.py +++ b/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') diff --git a/src/ocrmypdf/api.py b/src/ocrmypdf/api.py index a0cb19f5..a9fb047c 100644 --- a/src/ocrmypdf/api.py +++ b/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) diff --git a/src/ocrmypdf/helpers.py b/src/ocrmypdf/helpers.py index 964f7670..c55ac5f9 100644 --- a/src/ocrmypdf/helpers.py +++ b/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) diff --git a/src/ocrmypdf/hocrtransform.py b/src/ocrmypdf/hocrtransform.py index 6240d7ee..f1446535 100755 --- a/src/ocrmypdf/hocrtransform.py +++ b/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 diff --git a/src/ocrmypdf/pluginspec.py b/src/ocrmypdf/pluginspec.py index 3bab37e1..a96178f9 100644 --- a/src/ocrmypdf/pluginspec.py +++ b/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