Merge branch 'feature/lambda'

This commit is contained in:
James R. Barlow
2021-04-01 16:36:22 -07:00
21 changed files with 856 additions and 246 deletions
+1
View File
@@ -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
+4 -1
View File
@@ -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:
+118 -120
View File
@@ -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)
+8
View File
@@ -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])
+4
View File
@@ -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)
+15 -4
View File
@@ -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
+3 -1
View File
@@ -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)
+10 -9
View File
@@ -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")
+26 -13
View File
@@ -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:
+4
View File
@@ -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.
+169
View File
@@ -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))
@@ -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
+171
View File
@@ -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
+2 -2
View File
@@ -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
+76 -47
View File
@@ -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()
+55 -32
View File
@@ -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,
+125 -6
View File
@@ -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<firstresult>`.
"""
@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<firstresult>`.
"""
@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 <https://github.com/tqdm/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<firstresult>`.
"""
OrientationConfidence = namedtuple('OrientationConfidence', ('angle', 'confidence'))
"""Expresses an OCR engine's confidence in page rotation.