diff --git a/src/ocrmypdf/builtin_plugins/concurrency.py b/src/ocrmypdf/builtin_plugins/concurrency.py index 797087ae..1fe34d8c 100644 --- a/src/ocrmypdf/builtin_plugins/concurrency.py +++ b/src/ocrmypdf/builtin_plugins/concurrency.py @@ -20,9 +20,8 @@ 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, Union +from multiprocessing.pool import Pool, ThreadPool +from typing import Callable, Iterable, Optional, Tuple, Type, Union from tqdm import tqdm @@ -31,7 +30,10 @@ from ocrmypdf._logging import TqdmConsole from ocrmypdf.exceptions import InputFileError from ocrmypdf.helpers import remove_all_log_handlers +ProcessPool = Pool Queue = Union[multiprocessing.Queue, queue.Queue] +UserInit = Callable[[], None] +WorkerInit = Callable[[Queue, UserInit, int], None] def log_listener(q: Queue): @@ -62,7 +64,7 @@ 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): +def process_init(q: Queue, user_init: UserInit, loglevel) -> None: """Initialize a process pool worker""" # Ignore SIGINT (our parent process will kill us gracefully) @@ -85,7 +87,7 @@ def process_init(q: Queue, user_init: Callable[[], None], loglevel): return -def thread_init(_queue: Queue, user_init: Callable[[], None], _loglevel): +def thread_init(q: Queue, user_init: UserInit, loglevel) -> None: # As a thread, block SIGBUS so the main thread deals with it... with suppress(AttributeError): signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGBUS}) @@ -107,9 +109,9 @@ class StandardExecutor(Executor): task_finished: Callable, ): if use_threads: - log_queue = queue.Queue(-1) - pool_class = ThreadPool - initializer = thread_init + log_queue: Queue = queue.Queue(-1) + pool_class: Type[Pool] = ThreadPool + initializer: WorkerInit = thread_init else: log_queue = multiprocessing.Queue(-1) pool_class = ProcessPool diff --git a/src/ocrmypdf/extra_plugins/semfree.py b/src/ocrmypdf/extra_plugins/semfree.py index c84b1b2d..186206d9 100644 --- a/src/ocrmypdf/extra_plugins/semfree.py +++ b/src/ocrmypdf/extra_plugins/semfree.py @@ -28,7 +28,7 @@ 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, Iterator +from typing import Callable, Iterable, Iterator, List from ocrmypdf import Executor, hookimpl from ocrmypdf._concurrent import NullProgressBar @@ -60,7 +60,9 @@ def process_sigbus(*args): class ConnectionLogHandler(logging.handlers.QueueHandler): def __init__(self, conn: Connection) -> None: - super().__init__(None) + # sets the parent's queue to None - parent only touches queue + # in enqueue() which we override + super().__init__(None) # type: ignore self.conn = conn def enqueue(self, record): @@ -126,8 +128,8 @@ class LambdaExecutor(Executor): if not grouped_args: return - processes = [] - connections = [] + processes: List[Process] = [] + connections: List[Connection] = [] for chunk in grouped_args: parent_conn, child_conn = Pipe() @@ -152,6 +154,8 @@ class LambdaExecutor(Executor): with self.pbar_class(**tqdm_kwargs) as pbar: while connections: for r in wait(connections): + if not isinstance(r, Connection): + raise NotImplementedError("We only support Connection()") try: msg_type, msg = r.recv() except EOFError: