diff --git a/docs/installation.rst b/docs/installation.rst index 9dc46286..f3c952fc 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -163,7 +163,7 @@ To install ocrmypdf for the system: .. code-block:: bash - sudo pip3 install ocrmypdf + pip3 install ocrmypdf To install for the current user only: @@ -386,12 +386,16 @@ dependencies: To install ocrmypdf for the system: +.. code-block:: bash + # As root user pip3 install ocrmypdf ldconfig Or, to install for the current user only: +.. code-block:: bash + export PATH=$HOME/.local/bin:$PATH pip3 install --user ocrmypdf @@ -542,8 +546,9 @@ to change the PATH. .. warning:: As of early 2021, users have reported problems with the Microsoft Store version of - Python affected most third party Python packages including OCRmyPDF. Please use - Python downloaded from Python.org or Chocolatey as recommended here. + Python and OCRmyPDF. These issues affect many other third party Python packages. + Please download Python from Python.org or Chocolatey instead, and do not use the + Microsoft Store version. Windows Subsystem for Linux --------------------------- @@ -651,6 +656,21 @@ the latest version. However, PyPI and ``pip`` cannot address the fact that ``ocrmypdf`` depends on certain non-Python system libraries and programs being installed. +.. warning:: + + Debian and Ubuntu users: unfortunately, Debian and Ubuntu customize + Python in non-standard ways, and the nature of these customizations + varies from release to release. This can make for a frustrating + user experience. The instructions below work on almost all platforms that + have Python installed, except for Debian and Ubuntu, where you may need + to take additional steps. For best results on Debian and Ubuntu, use the + ``apt`` packages; or if these are too old, run + ``apt install python3-pip python3-venv``, create a virtual environment, + and install OCRmyPDF in that environment. + + `See here for more inforation on Debian-Python issues + `__. + For best results, first install `your platform's version `__ of ``ocrmypdf``, using the instructions elsewhere in this document. Then diff --git a/docs/plugins.rst b/docs/plugins.rst index e73523af..f0ae2d85 100644 --- a/docs/plugins.rst +++ b/docs/plugins.rst @@ -151,6 +151,18 @@ Custom command line arguments .. autofunction:: ocrmypdf.pluginspec.check_options +Execution and progress reporting +-------------------------------- + +.. autoclass: ocrmypdf.pluginspec.Executor + :members: + +.. autofunction:: ocrmypdf.pluginspec.get_logging_console + +.. autofunction:: ocrmypdf.pluginspec.get_executor + +.. autofunction:: ocrmypdf.pluginspec.get_progressbar_class + Applying special behavior before processing ------------------------------------------- diff --git a/src/ocrmypdf/__init__.py b/src/ocrmypdf/__init__.py index 081324e9..80586658 100644 --- a/src/ocrmypdf/__init__.py +++ b/src/ocrmypdf/__init__.py @@ -8,6 +8,7 @@ from pluggy import HookimplMarker as _HookimplMarker from ocrmypdf import helpers, hocrtransform, leptonica, pdfa, pdfinfo +from ocrmypdf._concurrent import Executor from ocrmypdf._jobcontext import PageContext, PdfContext from ocrmypdf._version import PROGRAM_NAME, __version__ from ocrmypdf.api import Verbosity, configure_logging, ocr diff --git a/src/ocrmypdf/__main__.py b/src/ocrmypdf/__main__.py index 3c897b72..1046a50c 100755 --- a/src/ocrmypdf/__main__.py +++ b/src/ocrmypdf/__main__.py @@ -47,7 +47,10 @@ def run(args=None): verbosity = Verbosity.quiet options.progress_bar = False configure_logging( - verbosity, progress_bar_friendly=options.progress_bar, manage_root_logger=True + verbosity, + progress_bar_friendly=options.progress_bar, + manage_root_logger=True, + plugin_manager=plugin_manager, ) log.debug('ocrmypdf %s', __version__) try: diff --git a/src/ocrmypdf/_concurrent.py b/src/ocrmypdf/_concurrent.py index 177ffe0a..5882d962 100644 --- a/src/ocrmypdf/_concurrent.py +++ b/src/ocrmypdf/_concurrent.py @@ -4,134 +4,132 @@ # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. - -import logging -import logging.handlers -import multiprocessing -import os -import signal import sys import threading -from contextlib import suppress -from multiprocessing import Pool as ProcessPool -from multiprocessing.dummy import Pool as ThreadPool +from abc import ABC, abstractmethod +from functools import partial from typing import Callable, Iterable, Optional -from tqdm import tqdm -from ocrmypdf.exceptions import InputFileError +def _task_noop(*_args, **_kwargs): + return -def log_listener(queue): - """Listen to the worker processes and forward the messages to logging +class NullProgressBar: + def __init__(self, **kwargs): + pass - For simplicity this is a thread rather than a process. Only one process - should actually write to sys.stderr or whatever we're using, so if this is - made into a process the main application needs to be directed to it. + def __enter__(self): + return self - See https://docs.python.org/3/howto/logging-cookbook.html#logging-to-a-single-file-from-multiple-processes + def __exit__(self, exc_type, exc_value, traceback): + return False + + def update(self, _arg=None): + return + + +class Executor(ABC): + pool_lock = threading.Lock() + pbar_class = NullProgressBar + + def __init__(self, *, pbar_class=None): + if pbar_class: + self.pbar_class = pbar_class + + def __call__( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Optional[Callable] = None, + task: Optional[Callable] = None, + task_arguments: Optional[Iterable] = None, + task_finished: Optional[Callable] = None, + ) -> None: + """ + Set up parallel execution and progress reporting. + + Args: + use_threads: If ``False``, the workload is the sort that will benefit from + running in a multiprocessing context (for example, it uses Python + heavily, and parallelizing it with threads is not expected to be + performant). + max_workers: The maximum number of workers that should be run. + tdqm_kwargs: Arguments to set up the progress bar. + worker_initializer: Called when a worker is initialized, in the worker's + execution context. If the child workers are processes, it must be + possible to marshall/pickle the worker initializer. + ``functools.partial`` can be used to bind parameters. + task: Called when the worker starts a new task, in the worker's execution + context. Must be possible to marshall to the worker. + task_finished: Called when a worker finishes a task, in the parent's + context. + task_arguments: An iterable that generates a group of parameters for each + task. This runs in the parent's context, but the parameters must be + marshallable to the worker. + """ + + if not task_arguments: + return # Nothing to do! + if not worker_initializer: + worker_initializer = _task_noop + if not task_finished: + task_finished = _task_noop + if not task: + task = _task_noop + + with self.pool_lock: + self._execute( + use_threads=use_threads, + max_workers=max_workers, + tqdm_kwargs=tqdm_kwargs, + worker_initializer=worker_initializer, + task=task, + task_arguments=task_arguments, + task_finished=task_finished, + ) + + @abstractmethod + def _execute( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + task_finished: Callable, + ): + """Custom executors should override this method.""" + + +def setup_executor(plugin_manager) -> Executor: + pbar_class = plugin_manager.hook.get_progressbar_class() + return plugin_manager.hook.get_executor(progressbar_class=pbar_class) + + +class SerialExecutor(Executor): + """Implements a purely sequential executor using the parallel protocol. + + The current process/thread will be the worker that executes all tasks + in order. As such, ``worker_initializer`` will never be called. """ - while True: - try: - record = queue.get() - if record is None: - break - logger = logging.getLogger(record.name) - logger.handle(record) - except Exception: # pylint: disable=broad-except - import traceback # pylint: disable=import-outside-toplevel - - print("Logging problem", file=sys.stderr) - traceback.print_exc(file=sys.stderr) - - -def process_sigbus(*args): - raise InputFileError("A worker process lost access to an input file") - - -def process_init(queue, user_init, loglevel): - """Initialize a process pool worker""" - - # Ignore SIGINT (our parent process will kill us gracefully) - signal.signal(signal.SIGINT, signal.SIG_IGN) - - # Install SIGBUS handler (so our parent process can abort somewhat gracefully) - with suppress(AttributeError): # Windows and Cygwin do not have SIGBUS - signal.signal(signal.SIGBUS, process_sigbus) - - # Reconfigure the root logger for this process to send all messages to a queue - h = logging.handlers.QueueHandler(queue) - root = logging.getLogger() - root.setLevel(loglevel) - root.handlers = [] - root.addHandler(h) - - if user_init: - user_init() - - -def thread_init(_queue, user_init, _loglevel): - # As a thread, block SIGBUS so the main thread deals with it... - with suppress(AttributeError): - # Windows and Cygwin do not have pthread_sigmask or SIGBUS - signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGBUS}) - if user_init: - user_init() - - -def exec_progress_pool( - *, - use_threads: bool, - max_workers: int, - tqdm_kwargs: dict, - task_initializer: Optional[Callable] = None, - task: Optional[Callable] = None, - task_arguments: Optional[Iterable] = None, - task_finished: Optional[Callable] = None, -): - log_queue: multiprocessing.Queue = multiprocessing.Queue(-1) - listener = threading.Thread(target=log_listener, args=(log_queue,)) - - if use_threads: - pool_class = ThreadPool - initializer = thread_init - else: - pool_class = ProcessPool - initializer = process_init - listener.start() - - with tqdm(**tqdm_kwargs) as pbar: - pool = pool_class( - processes=max_workers, - initializer=initializer, - initargs=(log_queue, task_initializer, logging.getLogger("").level), - ) - try: - results = pool.imap_unordered(task, task_arguments) - for result in results: - if task_finished: - task_finished(result, pbar) - else: - pbar.update() - except KeyboardInterrupt: - # Terminate pool so we exit instantly - pool.terminate() - # Don't try listener.join() here, will deadlock - raise - except Exception: - if not os.environ.get("PYTEST_CURRENT_TEST", ""): - # Unless inside pytest, exit immediately because no one wants - # to wait for child processes to finalize results that will be - # thrown away. Inside pytest, we want child processes to exit - # cleanly so that they output an error messages or coverage data - # we need from them. - pool.terminate() - raise - finally: - # Terminate log listener - log_queue.put_nowait(None) - pool.close() - pool.join() - - listener.join() + def _execute( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + task_finished: Callable, + ): + with self.pbar_class(**tqdm_kwargs) as pbar: + for args in task_arguments: + result = task(args) + task_finished(result, pbar) diff --git a/src/ocrmypdf/_exec/jbig2enc.py b/src/ocrmypdf/_exec/jbig2enc.py index 9311ec3f..2e8a058b 100644 --- a/src/ocrmypdf/_exec/jbig2enc.py +++ b/src/ocrmypdf/_exec/jbig2enc.py @@ -41,9 +41,17 @@ def convert_group(*, cwd, infiles, out_prefix): return proc +def convert_group_mp(args): + return convert_group(cwd=args[0], infiles=args[1], out_prefix=args[2]) + + def convert_single(*, cwd, infile, outfile): args = ['jbig2', '-p', infile] with open(outfile, 'wb') as fstdout: proc = run(args, cwd=cwd, stdout=fstdout, stderr=PIPE) proc.check_returncode() return proc + + +def convert_single_mp(args): + return convert_single(cwd=args[0], infile=args[1], outfile=args[2]) diff --git a/src/ocrmypdf/_exec/pngquant.py b/src/ocrmypdf/_exec/pngquant.py index 88a6370c..ca8a4542 100644 --- a/src/ocrmypdf/_exec/pngquant.py +++ b/src/ocrmypdf/_exec/pngquant.py @@ -59,3 +59,7 @@ def quantize(input_file: Path, output_file: Path, quality_min: int, quality_max: if result.returncode == 0: # input_file could be the same as output_file, so we defer the write output_file.write_bytes(result.stdout) + + +def quantize_mp(args): + return quantize(*args) diff --git a/src/ocrmypdf/_pipeline.py b/src/ocrmypdf/_pipeline.py index 747b2936..329a2cc0 100644 --- a/src/ocrmypdf/_pipeline.py +++ b/src/ocrmypdf/_pipeline.py @@ -19,9 +19,9 @@ import img2pdf import pikepdf from pikepdf.models.metadata import encode_pdf_date from PIL import Image, ImageColor, ImageDraw -from tqdm import tqdm from ocrmypdf import leptonica +from ocrmypdf._concurrent import Executor from ocrmypdf._exec import unpaper from ocrmypdf._jobcontext import PageContext, PdfContext from ocrmypdf._version import PROGRAM_NAME @@ -146,6 +146,8 @@ def triage(original_filename, input_file, output_file, options): def get_pdfinfo( input_file, + *, + executor: Executor, detailed_analysis=False, progbar=False, max_workers=None, @@ -158,6 +160,7 @@ def get_pdfinfo( progbar=progbar, max_workers=max_workers, check_pages=check_pages, + executor=executor, ) except pikepdf.PasswordError: raise EncryptedPdfError() @@ -611,6 +614,10 @@ def create_pdf_page_from_image( ) log.debug('convert done') + output_file = page_context.plugin_manager.hook.filter_pdf_page( + page=page_context, image_filename=image, output_pdf=output_file + ) + return output_file @@ -726,7 +733,11 @@ def convert_to_pdfa(input_pdf: Path, input_ps_stub: Path, context: PdfContext): output_file=output_file, compression=options.pdfa_image_compression, pdfa_part=options.output_type[-1], # is pdfa-1, pdfa-2, or pdfa-3 - progressbar_class=tqdm if options.progress_bar else None, + progressbar_class=( + context.plugin_manager.hook.get_progressbar_class() + if options.progress_bar + else None + ), ) return output_file @@ -812,13 +823,13 @@ def metadata_fixup(working_file: Path, context: PdfContext): return output_file -def optimize_pdf(input_file: Path, context: PdfContext): +def optimize_pdf(input_file: Path, context: PdfContext, executor: Executor): output_file = context.get_path('optimize.pdf') save_settings = dict( linearize=should_linearize(input_file, context), **get_pdf_save_settings(context.options.output_type), ) - optimize(input_file, output_file, context, save_settings) + optimize(input_file, output_file, context, save_settings, executor) return output_file diff --git a/src/ocrmypdf/_plugin_manager.py b/src/ocrmypdf/_plugin_manager.py index c7e77c07..3d245676 100644 --- a/src/ocrmypdf/_plugin_manager.py +++ b/src/ocrmypdf/_plugin_manager.py @@ -66,7 +66,9 @@ class OcrmypdfPluginManager(pluggy.PluginManager): # 1. Register builtins if self.__builtins: - for module in pkgutil.iter_modules(ocrmypdf.builtin_plugins.__path__): + for module in sorted( + pkgutil.iter_modules(ocrmypdf.builtin_plugins.__path__) + ): name = f'ocrmypdf.builtin_plugins.{module.name}' module = importlib.import_module(name) self.register(module) diff --git a/src/ocrmypdf/_sync.py b/src/ocrmypdf/_sync.py index 43e5008a..60e9192f 100644 --- a/src/ocrmypdf/_sync.py +++ b/src/ocrmypdf/_sync.py @@ -15,10 +15,9 @@ from pathlib import Path from tempfile import mkdtemp from typing import List, NamedTuple, Optional, Tuple -import pikepdf import PIL -from ocrmypdf._concurrent import exec_progress_pool +from ocrmypdf._concurrent import Executor, setup_executor from ocrmypdf._graft import OcrGrafter from ocrmypdf._jobcontext import PageContext, PdfContext, cleanup_working_files from ocrmypdf._logging import PageNumberFilter @@ -224,14 +223,14 @@ def exec_page_sync(page_context: PageContext): ) -def post_process(pdf_file, context: PdfContext): +def post_process(pdf_file, context: PdfContext, executor: Executor): pdf_out = pdf_file if context.options.output_type.startswith('pdfa'): ps_stub_out = generate_postscript_stub(context) pdf_out = convert_to_pdfa(pdf_out, ps_stub_out, context) pdf_out = metadata_fixup(pdf_out, context) - return optimize_pdf(pdf_out, context) + return optimize_pdf(pdf_out, context, executor) def worker_init(max_pixels: int): @@ -242,7 +241,7 @@ def worker_init(max_pixels: int): pikepdf_enable_mmap() -def exec_concurrent(context: PdfContext): +def exec_concurrent(context: PdfContext, executor: Executor): """Execute the pipeline concurrently""" # Run exec_page_sync on every page context @@ -269,7 +268,7 @@ def exec_concurrent(context: PdfContext): finally: tls.pageno = None - exec_progress_pool( + executor( use_threads=options.use_threads, max_workers=max_workers, tqdm_kwargs=dict( @@ -279,7 +278,7 @@ def exec_concurrent(context: PdfContext): unit_scale=0.5, disable=not options.progress_bar, ), - task_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS), + worker_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS), task=exec_page_sync, task_arguments=context.get_page_contexts(), task_finished=update_page, @@ -296,7 +295,7 @@ def exec_concurrent(context: PdfContext): # PDF/A and metadata log.info("Postprocessing...") - pdf = post_process(pdf, context) + pdf = post_process(pdf, context, executor) # Copy PDF file to destination copy_final(pdf, options.output_file, context) @@ -346,6 +345,7 @@ def run_pipeline(options, *, plugin_manager, api=False): pikepdf_enable_mmap() + executor = setup_executor(plugin_manager) try: check_requested_output_file(options) start_input_file, original_filename = create_input_file(options, work_folder) @@ -358,6 +358,7 @@ def run_pipeline(options, *, plugin_manager, api=False): # Gather pdfinfo and create context pdfinfo = get_pdfinfo( origin_pdf, + executor=executor, detailed_analysis=options.redo_ocr, progbar=options.progress_bar, max_workers=options.jobs if not options.use_threads else 1, # To help debug @@ -370,7 +371,7 @@ def run_pipeline(options, *, plugin_manager, api=False): validate_pdfinfo_options(context) # Execute the pipeline - exec_concurrent(context) + exec_concurrent(context, executor) if options.output_file == '-': log.info("Output sent to stdout") diff --git a/src/ocrmypdf/api.py b/src/ocrmypdf/api.py index 9a37cad6..9bce5352 100644 --- a/src/ocrmypdf/api.py +++ b/src/ocrmypdf/api.py @@ -15,7 +15,10 @@ from pathlib import Path from typing import AnyStr, BinaryIO, Iterable, Optional, Union from warnings import warn -from ocrmypdf._logging import PageNumberFilter, TqdmConsole +from ocrmypdf._logging import ( # pylint: disable=unused-import + PageNumberFilter, + TqdmConsole, +) from ocrmypdf._plugin_manager import get_plugin_manager from ocrmypdf._sync import run_pipeline from ocrmypdf._validation import check_options @@ -47,16 +50,20 @@ def configure_logging( verbosity: Verbosity, progress_bar_friendly: bool = True, manage_root_logger: bool = False, + plugin_manager=None, ): """Set up logging. Before calling :func:`ocrmypdf.ocr()`, you can use this function to - configure logging, if you want ocrmypdf's output to look like the ocrmypdf + configure logging if you want ocrmypdf's output to look like the ocrmypdf command line interface. It will register log handlers, log filters, and formatters, configure color logging to standard error, and adjust the log levels of third party libraries. Details of this are fine-tuned and subject to change. The ``verbosity`` argument is equivalent to the argument - ``--verbose`` and applies those settings. + ``--verbose`` and applies those settings. If you have a wrapper + script for ocrmypdf and you want it to be very similar to ocrmypdf, use this + function; if you are using ocrmypdf as part of an application that manages + its own logging, you probably do not want this function. If this function is not called, ocrmypdf will not configure logging, and it is up to the caller of ``ocrmypdf.ocr()`` to set up logging as it wishes using @@ -74,12 +81,11 @@ def configure_logging( their own debug logging. Args: - verbosity (Verbosity): Verbosity level. - progress_bar_friendly (bool): Install the TqdmConsole log handler, which is - compatible with the tqdm progress bar; without this log messages will - overwrite the progress bar - manage_root_logger (bool): Configure the process's root logger, to ensure - all log output is sent through + verbosity: Verbosity level. + progress_bar_friendly: If True (the default), install a custom log handler + that is compatible with progress bars and colored output. + manage_root_logger: Configure the process's root logger. + plugin_manager: The plugin manager, used for obtaining the custom log handler. Returns: The toplevel logger for ocrmypdf (or the root logger, if we are managing it). @@ -90,9 +96,11 @@ def configure_logging( log = logging.getLogger(prefix) log.setLevel(logging.DEBUG) - if progress_bar_friendly: - console = logging.StreamHandler(stream=TqdmConsole(sys.stderr)) - else: + console = None + if plugin_manager and progress_bar_friendly: + console = plugin_manager.hook.get_logging_console() + + if not console: console = logging.StreamHandler(stream=sys.stderr) if verbosity < 0: @@ -245,6 +253,7 @@ def ocr( # pylint: disable=unused-argument user_patterns: os.PathLike = None, fast_web_view: float = None, plugins: Iterable[StrPath] = None, + plugin_manager=None, keep_temporary_files: bool = None, progress_bar: bool = None, **kwargs, @@ -296,6 +305,9 @@ def ocr( # pylint: disable=unused-argument Returns: :class:`ocrmypdf.ExitCode` """ + if plugins and plugin_manager: + raise ValueError("plugins= and plugin_manager are mutually exclusive") + if not plugins: plugins = [] elif isinstance(plugins, (str, Path)): @@ -315,7 +327,8 @@ def ocr( # pylint: disable=unused-argument # they might install different plugins, and generally speaking we have areas # of code that use global state. - plugin_manager = get_plugin_manager(plugins) + if not plugin_manager: + plugin_manager = get_plugin_manager(plugins) plugin_manager.hook.add_options(parser=parser) # pylint: disable=no-member if 'verbose' in kwargs: diff --git a/src/ocrmypdf/builtin_plugins/__init__.py b/src/ocrmypdf/builtin_plugins/__init__.py index 61f9f6d6..05d8c70e 100644 --- a/src/ocrmypdf/builtin_plugins/__init__.py +++ b/src/ocrmypdf/builtin_plugins/__init__.py @@ -3,3 +3,7 @@ # This Source Code Form is subject to the terms of the Mozilla Public # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. + +# This file exists only mark builtin_plugins as a package. +# The plugin manager will not load it, so anything defined here may not be +# processed as a module. diff --git a/src/ocrmypdf/builtin_plugins/concurrency.py b/src/ocrmypdf/builtin_plugins/concurrency.py new file mode 100644 index 00000000..b58211e4 --- /dev/null +++ b/src/ocrmypdf/builtin_plugins/concurrency.py @@ -0,0 +1,169 @@ +# © 2020 James R. Barlow: github.com/jbarlow83 +# +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +# © 2020 James R. Barlow: github.com/jbarlow83 +# +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + + +import logging +import logging.handlers +import multiprocessing +import os +import queue +import signal +import sys +import threading +from contextlib import suppress +from multiprocessing import Pool as ProcessPool +from multiprocessing.pool import ThreadPool +from typing import Callable, Iterable, Optional, Union + +from tqdm import tqdm + +from ocrmypdf import Executor, hookimpl +from ocrmypdf._logging import TqdmConsole +from ocrmypdf.exceptions import InputFileError + +Queue = Union[multiprocessing.Queue, queue.Queue] + + +def log_listener(q: Queue): + """Listen to the worker processes and forward the messages to logging + + For simplicity this is a thread rather than a process. Only one process + should actually write to sys.stderr or whatever we're using, so if this is + made into a process the main application needs to be directed to it. + + See https://docs.python.org/3/howto/logging-cookbook.html#logging-to-a-single-file-from-multiple-processes + """ + + while True: + try: + record = q.get() + if record is None: + break + logger = logging.getLogger(record.name) + logger.handle(record) + except Exception: # pylint: disable=broad-except + import traceback # pylint: disable=import-outside-toplevel + + print("Logging problem", file=sys.stderr) + traceback.print_exc(file=sys.stderr) + + +def process_sigbus(*args): + raise InputFileError("A worker process lost access to an input file") + + +def process_init(q: Queue, user_init: Callable[[], None], loglevel): + """Initialize a process pool worker""" + + # Ignore SIGINT (our parent process will kill us gracefully) + signal.signal(signal.SIGINT, signal.SIG_IGN) + + # Install SIGBUS handler (so our parent process can abort somewhat gracefully) + with suppress(AttributeError): # Windows and Cygwin do not have SIGBUS + # Windows and Cygwin do not have pthread_sigmask or SIGBUS + signal.signal(signal.SIGBUS, process_sigbus) + + # Reconfigure the root logger for this process to send all messages to a queue + h = logging.handlers.QueueHandler(q) + root = logging.getLogger() + root.setLevel(loglevel) + root.handlers = [] + root.addHandler(h) + + user_init() + return + + +def thread_init(_queue: Queue, user_init: Callable[[], None], _loglevel): + # As a thread, block SIGBUS so the main thread deals with it... + with suppress(AttributeError): + signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGBUS}) + + user_init() + return + + +class StandardExecutor(Executor): + def _execute( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + task_finished: Callable, + ): + if use_threads: + log_queue = queue.Queue(-1) + pool_class = ThreadPool + initializer = thread_init + else: + log_queue = multiprocessing.Queue(-1) + pool_class = ProcessPool + initializer = process_init + + # Regardless of whether we use_threads for worker processes, the log_listener + # must be a thread + listener = threading.Thread(target=log_listener, args=(log_queue,)) + listener.start() + + with self.pbar_class(**tqdm_kwargs) as pbar: + pool = pool_class( + processes=max_workers, + initializer=initializer, + initargs=(log_queue, worker_initializer, logging.getLogger("").level), + ) + try: + results = pool.imap_unordered(task, task_arguments) + for result in results: + if task_finished: + task_finished(result, pbar) + else: + pbar.update() + except KeyboardInterrupt: + # Terminate pool so we exit instantly + pool.terminate() + # Don't try listener.join() here, will deadlock + raise + except Exception: + if not os.environ.get("PYTEST_CURRENT_TEST", ""): + # Unless inside pytest, exit immediately because no one wants + # to wait for child processes to finalize results that will be + # thrown away. Inside pytest, we want child processes to exit + # cleanly so that they output an error messages or coverage data + # we need from them. + pool.terminate() + raise + finally: + # Terminate log listener + log_queue.put_nowait(None) + pool.close() + pool.join() + + listener.join() + + +@hookimpl +def get_executor(progressbar_class): + return StandardExecutor(pbar_class=progressbar_class) + + +@hookimpl +def get_progressbar_class(): + return tqdm + + +@hookimpl +def get_logging_console(): + return logging.StreamHandler(stream=TqdmConsole(sys.stderr)) diff --git a/src/ocrmypdf/builtin_plugins/default_filters.py b/src/ocrmypdf/builtin_plugins/default_filters.py new file mode 100644 index 00000000..83a8f3fc --- /dev/null +++ b/src/ocrmypdf/builtin_plugins/default_filters.py @@ -0,0 +1,12 @@ +# © 2021 James R. Barlow: github.com/jbarlow83 +# +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +from ocrmypdf import hookimpl + + +@hookimpl +def filter_pdf_page(page, image_filename, output_pdf): + return output_pdf diff --git a/src/ocrmypdf/extra_plugins/__init__.py b/src/ocrmypdf/extra_plugins/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/ocrmypdf/extra_plugins/awslambda.py b/src/ocrmypdf/extra_plugins/awslambda.py new file mode 100644 index 00000000..23fe9799 --- /dev/null +++ b/src/ocrmypdf/extra_plugins/awslambda.py @@ -0,0 +1,171 @@ +# © 2021 James R. Barlow: github.com/jbarlow83 +# +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + + +"""Alternate executor to support OCRmyPDF in AWS Lambda""" + + +import logging +import logging.handlers +import signal +from contextlib import suppress +from enum import Enum, auto +from itertools import islice, repeat, takewhile, zip_longest +from multiprocessing import Pipe, Process +from multiprocessing.connection import Connection, wait +from typing import Callable, Iterable, Optional +from unittest.mock import Mock + +from ocrmypdf import Executor, hookimpl +from ocrmypdf._concurrent import NullProgressBar +from ocrmypdf.exceptions import InputFileError + + +class MessageType(Enum): + exception = auto() + result = auto() + complete = auto() + + +def split_every(n: int, iterable: Iterable): + iterator = iter(iterable) + return takewhile(bool, (list(islice(iterator, n)) for _ in repeat(None))) + + +def process_sigbus(*args): + raise InputFileError("A worker process lost access to an input file") + + +class ConnectionLogHandler(logging.handlers.QueueHandler): + def __init__(self, conn: Connection) -> None: + super().__init__(None) + self.conn = conn + + def enqueue(self, record): + self.conn.send(('log', record)) + + +def process_loop( + conn: Connection, user_init: Callable[[], None], loglevel, task, task_args +): + """Initialize a process pool worker""" + + # Install SIGBUS handler (so our parent process can abort somewhat gracefully) + with suppress(AttributeError): # Windows and Cygwin do not have SIGBUS + # Windows and Cygwin do not have pthread_sigmask or SIGBUS + signal.signal(signal.SIGBUS, process_sigbus) + + # Reconfigure the root logger for this process to send all messages to a queue + h = ConnectionLogHandler(conn) + root = logging.getLogger() + root.setLevel(loglevel) + root.handlers = [] + root.addHandler(h) + + user_init() + + for args in task_args: + try: + result = task(args) + except Exception as e: + conn.send((MessageType.exception, e)) + break + else: + conn.send((MessageType.result, result)) + + conn.send((MessageType.complete, None)) + conn.close() + return + + +class LambdaExecutor(Executor): + def _execute( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + task_finished: Callable, + ): + if use_threads and max_workers == 1: + for args in task_arguments: + result = task(args) + task_finished(result, self.pbar_class) + return + + task_arguments = list(task_arguments) + grouped_args = list( + zip_longest(*list(split_every(max_workers, task_arguments))) + ) + if not grouped_args: + return + + processes = [] + connections = [] + for chunk in grouped_args: + parent_conn, child_conn = Pipe() + + worker_args = [args for args in chunk if args is not None] + process = Process( + target=process_loop, + args=( + child_conn, + worker_initializer, + logging.getLogger("").level, + task, + worker_args, + ), + ) + process.daemon = True + processes.append(process) + connections.append(parent_conn) + + for process in processes: + process.start() + + with self.pbar_class(**tqdm_kwargs) as pbar: + while connections: + for r in wait(connections): + try: + msg_type, msg = r.recv() + except EOFError: + connections.remove(r) + continue + + if msg_type == MessageType.result: + if task_finished: + task_finished(msg, pbar) + elif msg_type == 'log': + record = msg + logger = logging.getLogger(record.name) + logger.handle(record) + elif msg_type == MessageType.complete: + connections.remove(r) + elif msg_type == MessageType.exception: + for process in processes: + process.terminate() + raise msg + + for process in processes: + process.join() + + +@hookimpl +def get_executor(progressbar_class): + return LambdaExecutor(pbar_class=progressbar_class) + + +@hookimpl +def get_logging_console(): + return logging.StreamHandler() + + +@hookimpl +def get_progressbar_class(): + return NullProgressBar diff --git a/src/ocrmypdf/leptonica.py b/src/ocrmypdf/leptonica.py index 91715939..c8759cfa 100644 --- a/src/ocrmypdf/leptonica.py +++ b/src/ocrmypdf/leptonica.py @@ -15,7 +15,6 @@ import logging import os import sys import threading -import warnings from collections import deque from collections.abc import Sequence from contextlib import suppress @@ -24,6 +23,7 @@ from functools import lru_cache from io import BytesIO, UnsupportedOperation from os import fspath from tempfile import TemporaryFile +from warnings import warn from ocrmypdf.exceptions import MissingDependencyError from ocrmypdf.lib._leptonica import ffi @@ -389,7 +389,7 @@ class Pix(LeptonicaObject): @classmethod def read(cls, path): - warnings.warn('Use Pix.open() instead', DeprecationWarning) + warn('Use Pix.open() instead', DeprecationWarning) return cls.open(path) @classmethod diff --git a/src/ocrmypdf/optimize.py b/src/ocrmypdf/optimize.py index b420948d..a26057f4 100644 --- a/src/ocrmypdf/optimize.py +++ b/src/ocrmypdf/optimize.py @@ -9,7 +9,6 @@ import logging import sys import tempfile from collections import defaultdict -from functools import partial from os import fspath from pathlib import Path from typing import ( @@ -29,10 +28,9 @@ import img2pdf import pikepdf from pikepdf import Dictionary, Name, Object, Pdf, PdfImage from PIL import Image -from tqdm import tqdm from ocrmypdf import leptonica -from ocrmypdf._concurrent import exec_progress_pool +from ocrmypdf._concurrent import Executor, SerialExecutor from ocrmypdf._exec import jbig2enc, pngquant from ocrmypdf._jobcontext import PdfContext from ocrmypdf.exceptions import OutputFileAccessError @@ -301,17 +299,17 @@ def extract_images_jbig2(pike: Pdf, root: Path, options) -> Dict[int, List[XrefE def _produce_jbig2_images( - jbig2_groups: Dict[int, List[XrefExt]], root: Path, options + jbig2_groups: Dict[int, List[XrefExt]], root: Path, options, executor: Executor ) -> None: """Produce JBIG2 images from their groups""" def jbig2_group_args(root: Path, groups: Dict[int, List[XrefExt]]): for group, xref_exts in groups.items(): prefix = f'group{group:08d}' - yield dict( - cwd=fspath(root), - infiles=(img_name(root, xref, ext) for xref, ext in xref_exts), - out_prefix=prefix, + yield ( + fspath(root), # =cwd + (img_name(root, xref, ext) for xref, ext in xref_exts), # =infiles + prefix, # =out_prefix ) def jbig2_single_args(root, groups: Dict[int, List[XrefExt]]): @@ -320,23 +318,20 @@ def _produce_jbig2_images( # Second loop is to ensure multiple images per page are unpacked for n, xref_ext in enumerate(xref_exts): xref, ext = xref_ext - yield dict( - cwd=fspath(root), - infile=img_name(root, xref, ext), - outfile=root / f'{prefix}.{n:04d}', + yield ( + fspath(root), + img_name(root, xref, ext), + root / f'{prefix}.{n:04d}', ) - def convert_generic(fn, kwargs_dict): - return fn(**kwargs_dict) - if options.jbig2_page_group_size > 1: jbig2_args = jbig2_group_args - jbig2_convert = partial(convert_generic, jbig2enc.convert_group) + jbig2_convert = jbig2enc.convert_group_mp else: jbig2_args = jbig2_single_args - jbig2_convert = partial(convert_generic, jbig2enc.convert_single) + jbig2_convert = jbig2enc.convert_single_mp - exec_progress_pool( + executor( use_threads=True, max_workers=options.jobs, tqdm_kwargs=dict( @@ -351,7 +346,11 @@ def _produce_jbig2_images( def convert_to_jbig2( - pike: Pdf, jbig2_groups: Dict[int, List[XrefExt]], root: Path, options + pike: Pdf, + jbig2_groups: Dict[int, List[XrefExt]], + root: Path, + options, + executor: Executor, ) -> None: """Convert images to JBIG2 and insert into PDF. @@ -366,7 +365,7 @@ def convert_to_jbig2( and needs no dictionary. Currently this must be lossless JBIG2. """ - _produce_jbig2_images(jbig2_groups, root, options) + _produce_jbig2_images(jbig2_groups, root, options, executor) for group, xref_exts in jbig2_groups.items(): prefix = f'group{group:08d}' @@ -390,27 +389,53 @@ def convert_to_jbig2( ) -def transcode_jpegs(pike: Pdf, jpegs: Sequence[Xref], root: Path, options) -> None: - for xref in tqdm( - jpegs, desc="JPEGs", unit='image', disable=not options.progress_bar - ): - in_jpg = jpg_name(root, xref) - opt_jpg = in_jpg.with_suffix('.opt.jpg') +def _optimize_jpeg(args): + xref, in_jpg, opt_jpg, jpeg_quality = args - # This produces a debug warning from PIL - # DEBUG:PIL.Image:Error closing: 'NoneType' object has no attribute - # 'close'. Seems to be mostly harmless - # https://github.com/python-pillow/Pillow/issues/1144 - with Image.open(in_jpg) as im: - im.save(opt_jpg, optimize=True, quality=options.jpeg_quality) + # This may produce a debug warning from PIL + # DEBUG:PIL.Image:Error closing: 'NoneType' object has no attribute + # 'close'. Seems to be mostly harmless + # https://github.com/python-pillow/Pillow/issues/1144 + with Image.open(in_jpg) as im: + im.save(opt_jpg, optimize=True, quality=jpeg_quality) - if opt_jpg.stat().st_size > in_jpg.stat().st_size: - log.debug("xref %s, jpeg, made larger - skip", xref) - continue + if opt_jpg.stat().st_size > in_jpg.stat().st_size: + log.debug("xref %s, jpeg, made larger - skip", xref) + opt_jpg.unlink() + opt_jpg = None + return xref, opt_jpg - compdata = leptonica.CompressedData.open(opt_jpg) - im_obj = pike.get_object(xref, 0) - im_obj.write(compdata.read(), filter=Name.DCTDecode) + +def transcode_jpegs( + pike: Pdf, jpegs: Sequence[Xref], root: Path, options, executor +) -> None: + def jpeg_args(): + for xref in jpegs: + in_jpg = jpg_name(root, xref) + opt_jpg = in_jpg.with_suffix('.opt.jpg') + yield xref, in_jpg, opt_jpg, options.jpeg_quality + + def finish_jpeg(result, pbar): + xref, opt_jpg = result + if opt_jpg: + compdata = leptonica.CompressedData.open(opt_jpg) + im_obj = pike.get_object(xref, 0) + im_obj.write(compdata.read(), filter=Name.DCTDecode) + pbar.update() + + executor( + use_threads=True, # Processes are significantly slower at this task + max_workers=options.jobs, + tqdm_kwargs=dict( + desc="JPEGs", + total=len(jpegs), + unit='image', + disable=not options.progress_bar, + ), + task=_optimize_jpeg, + task_arguments=jpeg_args(), + task_finished=finish_jpeg, + ) def _transcode_png(pike: Pdf, filename: Path, xref: Xref) -> bool: @@ -460,6 +485,7 @@ def transcode_pngs( image_name_fn: Callable[[Path, Xref], Path], root: Path, options, + executor, ) -> None: modified: MutableSet[Xref] = set() if options.optimize >= 2: @@ -479,10 +505,7 @@ def transcode_pngs( ) modified.add(xref) - def pngquant_fn(args): - pngquant.quantize(*args) - - exec_progress_pool( + executor( use_threads=True, max_workers=options.jobs, tqdm_kwargs=dict( @@ -491,7 +514,7 @@ def transcode_pngs( unit='image', disable=not options.progress_bar, ), - task=pngquant_fn, + task=pngquant.quantize_mp, task_arguments=pngquant_args(), ) @@ -575,7 +598,13 @@ def rewrite_png(pike: Pdf, im_obj: Object, compdata) -> None: # pragma: no cove im_obj.write(compdata.read(), filter=Name.FlateDecode, decode_parms=dparms) -def optimize(input_file: Path, output_file: Path, context, save_settings) -> None: +def optimize( + input_file: Path, + output_file: Path, + context, + save_settings, + executor: Executor = SerialExecutor(), +) -> None: options = context.options if options.optimize == 0: safe_symlink(input_file, output_file) @@ -593,14 +622,14 @@ def optimize(input_file: Path, output_file: Path, context, save_settings) -> Non root.mkdir(exist_ok=True) jpegs, pngs = extract_images_generic(pike, root, options) - transcode_jpegs(pike, jpegs, root, options) + transcode_jpegs(pike, jpegs, root, options, executor) # if options.optimize >= 2: # Try pngifying the jpegs # transcode_pngs(pike, jpegs, jpg_name, root, options) - transcode_pngs(pike, pngs, png_name, root, options) + transcode_pngs(pike, pngs, png_name, root, options, executor) jbig2_groups = extract_images_jbig2(pike, root, options) - convert_to_jbig2(pike, jbig2_groups, root, options) + convert_to_jbig2(pike, jbig2_groups, root, options, executor) target_file = output_file.with_suffix('.opt.pdf') pike.remove_unreferenced_resources() diff --git a/src/ocrmypdf/pdfinfo/info.py b/src/ocrmypdf/pdfinfo/info.py index 0ee921b0..50c77b0b 100644 --- a/src/ocrmypdf/pdfinfo/info.py +++ b/src/ocrmypdf/pdfinfo/info.py @@ -6,9 +6,11 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. +import atexit import logging import re from collections import defaultdict, namedtuple +from contextlib import ExitStack from decimal import Decimal from enum import Enum from functools import partial @@ -21,7 +23,7 @@ from warnings import warn import pikepdf from pikepdf import Object, Pdf, PdfMatrix -from ocrmypdf._concurrent import exec_progress_pool +from ocrmypdf._concurrent import Executor, SerialExecutor from ocrmypdf.exceptions import EncryptedPdfError, InputFileError from ocrmypdf.helpers import Resolution, available_cpu_count, pikepdf_enable_mmap from ocrmypdf.pdfinfo.layout import get_page_analysis, get_text_boxes @@ -571,29 +573,42 @@ def simplify_textboxes(miner, textbox_getter) -> Iterator[TextboxInfo]: worker_pdf = None -def _pdf_pageinfo_sync_init(infile: Path, pdfminer_loglevel): +def _pdf_pageinfo_sync_init(pdf: Pdf, infile: Path, pdfminer_loglevel): global worker_pdf # pylint: disable=global-statement pikepdf_enable_mmap() logging.getLogger('pdfminer').setLevel(pdfminer_loglevel) - # If this function is called as a thread initializer, we need a messy hack - # to close worker_pdf. If called as a process, it will be released when the - # process is terminated. - worker_pdf = pikepdf.open(infile) + # If the pdf is not opened, open a copy for our worker process to use + if pdf is None: + worker_pdf = pikepdf.open(infile) + + def on_process_close(): + worker_pdf.close() + + # Close when this process exits + atexit.register(on_process_close) def _pdf_pageinfo_sync(args): - global worker_pdf # pylint: disable=global-statement - pageno, infile, check_pages, detailed_analysis = args - page = PageInfo(worker_pdf, pageno, infile, check_pages, detailed_analysis) - return page + pageno, thread_pdf, infile, check_pages, detailed_analysis = args + pdf = thread_pdf if thread_pdf is not None else worker_pdf + with ExitStack() as stack: + if not pdf: # When called with SerialExecutor + pdf = stack.enter_context(pikepdf.open(infile)) + page = PageInfo(pdf, pageno, infile, check_pages, detailed_analysis) + return page def _pdf_pageinfo_concurrent( - pdf, infile, progbar, max_workers, check_pages, detailed_analysis=False + pdf, + executor: Executor, + infile, + progbar, + max_workers, + check_pages, + detailed_analysis=False, ): - global worker_pdf # pylint: disable=global-statement pages = [None] * len(pdf.pages) def update_pageinfo(result, pbar): @@ -607,7 +622,6 @@ def _pdf_pageinfo_concurrent( max_workers = available_cpu_count() total = len(pdf.pages) - contexts = ((n, infile, check_pages, detailed_analysis) for n in range(total)) use_threads = False # No performance gain if threaded due to GIL n_workers = min(1 + len(pages) // 4, max_workers) @@ -616,25 +630,31 @@ def _pdf_pageinfo_concurrent( # a separate process. use_threads = True - try: - exec_progress_pool( - use_threads=use_threads, - max_workers=n_workers, - tqdm_kwargs=dict( - total=total, desc="Scanning contents", unit='page', disable=not progbar - ), - task_initializer=partial( - _pdf_pageinfo_sync_init, infile, logging.getLogger('pdfminer').level - ), - task=_pdf_pageinfo_sync, - task_arguments=contexts, - task_finished=update_pageinfo, - ) - finally: - if worker_pdf and use_threads: - assert n_workers == 1, "Should have only one worker when threaded" - # This is messy, but if we ran in thread, close worker_pdf - worker_pdf.close() + # If we use a thread, we can pass the already-open Pdf for them to use + # If we use processes, we pass a None which tells the init function to open its + # own + initial_pdf = pdf if use_threads else None + + contexts = ( + (n, initial_pdf, infile, check_pages, detailed_analysis) for n in range(total) + ) + assert n_workers == 1 if use_threads else n_workers >= 1, "Not multithreadable" + executor( + use_threads=use_threads, + max_workers=n_workers, + tqdm_kwargs=dict( + total=total, desc="Scanning contents", unit='page', disable=not progbar + ), + worker_initializer=partial( + _pdf_pageinfo_sync_init, + initial_pdf, + infile, + logging.getLogger('pdfminer').level, + ), + task=_pdf_pageinfo_sync, + task_arguments=contexts, + task_finished=update_pageinfo, + ) return pages @@ -819,10 +839,12 @@ class PdfInfo: def __init__( self, infile, + *, detailed_analysis: bool = False, progbar: bool = False, max_workers: int = None, check_pages=None, + executor: Executor = SerialExecutor(), ): self._infile = infile if check_pages is None: @@ -833,6 +855,7 @@ class PdfInfo: raise EncryptedPdfError() # Triggered by encryption with empty passwd self._pages = _pdf_pageinfo_concurrent( pdf, + executor, infile, progbar, max_workers, diff --git a/src/ocrmypdf/pluginspec.py b/src/ocrmypdf/pluginspec.py index 9aee6596..5fa0b3ca 100644 --- a/src/ocrmypdf/pluginspec.py +++ b/src/ocrmypdf/pluginspec.py @@ -8,11 +8,13 @@ from abc import ABC, abstractmethod, abstractstaticmethod from argparse import ArgumentParser, Namespace from collections import namedtuple +from logging import Handler from pathlib import Path -from typing import TYPE_CHECKING, AbstractSet, List, Optional +from typing import TYPE_CHECKING, AbstractSet, Callable, Iterable, List, Optional import pluggy +from ocrmypdf._concurrent import Executor from ocrmypdf.helpers import Resolution if TYPE_CHECKING: @@ -26,6 +28,18 @@ hookspec = pluggy.HookspecMarker('ocrmypdf') # pylint: disable=unused-argument +@hookspec(firstresult=True) +def get_logging_console() -> Handler: + """Returns a custom logging handler. + + Generally this is necessary when both logging output and a progress bar are both + outputting to ``sys.stderr``. + + Note: + This is a :ref:`firstresult hook`. + """ + + @hookspec def add_options(parser: ArgumentParser) -> None: """Allows the plugin to add its own command line and API arguments. @@ -62,6 +76,67 @@ def check_options(options: Namespace) -> None: """ +@hookspec(firstresult=True) +def get_executor(progressbar_class) -> Executor: + """Called to obtain an object that manages parallel execution. + + This may be used to replace OCRmyPDF's default parallel execution system + with a third party alternative. For example, you could make OCRmyPDF run in a + distributed environment. + + OCRmyPDF's executors are analogous to the standard Python executors in + ``conconcurrent.futures``, but they do not work the same way. Executors may + be reused for different, unrelated batch operations, since all of the context + for a given job are passed to :meth:`Executor.__call__`. + + Should be of type :class:`Executor` or otherwise conforming to the protocol + of that call. + + Arguments: + progressbar_class: A progress bar class, which will be created when + + Note: + This hook will be called from the main process, and may modify global state + before child worker processes are forked. + Note: + This is a :ref:`firstresult hook`. + """ + + +@hookspec(firstresult=True) +def get_progressbar_class(): + """Called to obtain a class that can be used to monitor progress. + + A progress bar is assumed, but this could be used for any type of monitoring. + + The class should follow a tqdm-like protocol. Calling the class should return + a new progress bar object, which is activated with ``__enter__`` and terminated + ``__exit__``. An update method is called whenever the progress bar is updated. + Progress bar objects will not be reused; a new one will be created for each + group of tasks. + + The progress bar is held in the main process/thread and not updated by child + process/threads. When a child notifies the parent of completed work, the + parent updates the progress bar. + + The arguments are the same as `tqdm `_ accepts. + + Progress bars should never write to ``sys.stdout``, or they will corrupt the + output if OCRmyPDF writes a PDF to standard output. + + The type of events that OCRmyPDF reports to a progress bar may change in + minor releases. + + Here is how OCRmyPDF will use the progress bar: + + Example: + pbar_class = pm.hook.get_progressbar_class() + with pbar_class(**tqdm_kwargs) as pbar: + ... + pbar.update(1) + """ + + @hookspec def validate(pdfinfo: 'PdfInfo', options: Namespace) -> None: """Called to give a plugin an opportunity to review *options* and *pdfinfo*. @@ -163,11 +238,11 @@ def filter_page_image(page: 'PageContext', image_filename: Path) -> Path: will be resized and the OCR layer misaligned. OCRmyPDF does not nothing to enforce these constraints; it is up to the plugin to do sensible things. - OCRmyPDF will create the PDF page based on the image format used. If you - convert the image to a JPEG, the output page will be created as a JPEG, etc. - If you change the colorspace, that change will be kept. Note that the - OCRmyPDF image optimization stage, if enabled, may ultimately chose a - different format. + OCRmyPDF will create the PDF page based on the image format used (unless the + hook is overriden). If you convert the image to a JPEG, the output page will + be created as a JPEG, etc. If you change the colorspace, that change will be + kept. Note that the OCRmyPDF image optimization stage, if enabled, may + ultimately chose a different format. If the return value is a file that does not exist, ``FileNotFoundError`` will occur. The return value should be a path to a file in the same folder @@ -185,6 +260,50 @@ def filter_page_image(page: 'PageContext', image_filename: Path) -> Path: """ +@hookspec(firstresult=True) +def filter_pdf_page( + page: 'PageContext', image_filename: Path, output_pdf: Path +) -> Path: + """Called to convert a filtered whole page image into a PDF. + + A whole page image is only produced when preprocessing command line arguments + are issued or when ``--force-ocr`` is issued. If no whole page is image is + produced for a given page, this function will not be called. This is not + the image that will be shown to OCR. The whole page image is filtered in + the hook above, ``filter_page_image``, then this function is called for + PDF conversion. + + This function will only be called when OCRmyPDF runs in a mode such as + "force OCR" mode where rasterizing of all content is performed. + + Clever things could be done at this stage such as segmenting the page image into + color regions or vector equivalents. + + The provider of the hook implementation is responsible for ensuring that the + OCR text layer is aligned with the PDF produced here, or text misalignment + will result. + + Currently this function must produce a single page PDF or the pipeline will + fail. If the intent is to remove the PDF, then create a single page empty + PDF. + + Args: + page: Context for this page. + image_filename: Filename of the input image used to create output_pdf, + for "reference" if recreating the output_pdf entirely. + output_pdf: The previous created output_pdf. + + Returns: + output_pdf + + Note: + This hook will be called from child processes. Modifying global state + will not affect the main process or other child processes. + Note: + This is a :ref:`firstresult hook`. + """ + + OrientationConfidence = namedtuple('OrientationConfidence', ('angle', 'confidence')) """Expresses an OCR engine's confidence in page rotation. diff --git a/tests/test_validation.py b/tests/test_validation.py index 3393fd9f..7b2f1ead 100644 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -6,12 +6,13 @@ import logging -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pikepdf import pytest from ocrmypdf import _validation as vd +from ocrmypdf._concurrent import NullProgressBar, SerialExecutor from ocrmypdf._plugin_manager import get_plugin_manager from ocrmypdf.api import create_options from ocrmypdf.cli import get_parser @@ -173,13 +174,22 @@ def test_false_action_store_true(): def test_no_progress_bar(progress_bar, resources): opts = make_opts(progress_bar=progress_bar, input_file=(resources / 'trivial.pdf')) plugin_manager = get_plugin_manager(opts.plugins) - with patch('ocrmypdf._concurrent.tqdm', autospec=True) as tqdmpatch: - vd._check_options(opts, plugin_manager, set()) - pdfinfo = PdfInfo(opts.input_file, progbar=opts.progress_bar) - assert pdfinfo is not None - assert tqdmpatch.called - _args, kwargs = tqdmpatch.call_args - assert kwargs['disable'] != progress_bar + + vd._check_options(opts, plugin_manager, set()) + + pbar_disabled = None + + class CheckProgressBar(NullProgressBar): + def __init__(self, disable, **kwargs): + nonlocal pbar_disabled + pbar_disabled = disable + super().__init__(disable=disable, **kwargs) + + executor = SerialExecutor(pbar_class=CheckProgressBar) + pdfinfo = PdfInfo(opts.input_file, progbar=opts.progress_bar, executor=executor) + + assert pdfinfo is not None + assert pbar_disabled is not None and pbar_disabled != progress_bar def test_language_warning(caplog):