Refactor to eliminate global state in _concurrent

src/ocrmypdf/extra_plugins/awslambda.py

@@ -19,7 +19,7 @@ from multiprocessing.connection import Connection, wait
 from typing import Callable, Iterable, Optional
 from unittest.mock import Mock
 
-from ocrmypdf import hookimpl
+from ocrmypdf import Executor, hookimpl
 from ocrmypdf.exceptions import InputFileError
 
 
@@ -80,97 +80,101 @@ def process_loop(
     return
 
 
-def lambda_pool_impl(
-    *,
-    use_threads: bool,
-    max_workers: int,
-    tqdm_kwargs: dict,
-    worker_initializer: Callable,
-    task: Callable,
-    task_arguments: Optional[Iterable] = None,
-    task_finished: Callable,
-):
-    pbar = Mock()
+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,
+    ):
+        pbar = Mock()
 
-    if use_threads and max_workers == 1:
-        for args in task_arguments:
-            result = task(args)
-            task_finished(result, pbar)
-        return
+        if use_threads and max_workers == 1:
+            for args in task_arguments:
+                result = task(args)
+                task_finished(result, pbar)
+            return
 
-    _lambda_pool_impl(
-        max_workers=max_workers,
-        worker_initializer=worker_initializer,
-        task=task,
-        task_arguments=task_arguments,
-        task_finished=task_finished,
-        pbar=pbar,
-    )
-
-
-def _lambda_pool_impl(
-    *,
-    max_workers: int,
-    worker_initializer: Callable,
-    task: Callable,
-    task_arguments: Optional[Iterable] = None,
-    task_finished: Callable,
-    pbar,
-):
-    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,
-            ),
+        self._lambda_pool_impl(
+            max_workers=max_workers,
+            worker_initializer=worker_initializer,
+            task=task,
+            task_arguments=task_arguments,
+            task_finished=task_finished,
+            pbar=pbar,
         )
-        process.daemon = True
-        processes.append(process)
-        connections.append(parent_conn)
 
-    for process in processes:
-        process.start()
+    def _lambda_pool_impl(
+        self,
+        *,
+        max_workers: int,
+        worker_initializer: Callable,
+        task: Callable,
+        task_arguments: Iterable,
+        task_finished: Callable,
+        pbar,
+    ):
+        task_arguments = list(task_arguments)
+        grouped_args = list(
+            zip_longest(*list(split_every(max_workers, task_arguments)))
+        )
+        if not grouped_args:
+            return
 
-    while connections:
-        for r in wait(connections):
-            try:
-                msg_type, msg = r.recv()
-            except EOFError:
-                connections.remove(r)
-                continue
+        processes = []
+        connections = []
+        for chunk in grouped_args:
+            parent_conn, child_conn = Pipe()
 
-            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
+            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.join()
+        for process in processes:
+            process.start()
+
+        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_parallel_executor():
-    return lambda_pool_impl
+def get_executor():
+    return LambdaExecutor()