Refactor - decouple progressbar from executor

docs/plugins.rst

@@ -130,6 +130,16 @@ Custom command line arguments
 
 .. autofunction:: ocrmypdf.pluginspec.check_options
 
+Execution and progress reporting
+--------------------------------
+
+.. autoclass: ocrmypdf.pluginspec.Executor
+    :members:
+
+.. autofunction:: ocrmypdf.pluginspec.get_executor
+
+.. autofunction:: ocrmypdf.pluginspec.get_progress_bar
+
 Applying special behavior before processing
 -------------------------------------------
 

src/ocrmypdf/_concurrent.py

@@ -10,15 +10,32 @@ from abc import ABC, abstractmethod
 from functools import partial
 from typing import Callable, Iterable, Optional
 
-from tqdm import tqdm
-
 
 def _task_noop(*_args, **_kwargs):
     return
 
 
+class NullProgressBar:
+    def __init__(self, **kwargs):
+        pass
+
+    def __enter__(self):
+        return self
+
+    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,
@@ -32,17 +49,21 @@ class Executor(ABC):
         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
+            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 the worker is initialized, in the worker's
-                execution context. Must be possible to marshall to the worker.
+            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 marshallable to the worker.
+                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
@@ -86,7 +107,8 @@ class Executor(ABC):
 
 
 def setup_executor(plugin_manager) -> Executor:
-    return plugin_manager.hook.get_executor()
+    pbar_class = plugin_manager.hook.get_progress_bar()
+    return plugin_manager.hook.get_executor(pbar_class=pbar_class)
 
 
 class SerialExecutor(Executor):
@@ -107,7 +129,7 @@ class SerialExecutor(Executor):
         task_arguments: Iterable,
         task_finished: Callable,
     ):
-        with tqdm(**tqdm_kwargs) as pbar:
+        with self.pbar_class(**tqdm_kwargs) as pbar:
             for args in task_arguments:
                 result = task(args)
                 task_finished(result, pbar)

src/ocrmypdf/builtin_plugins/concurrency.py

@@ -117,7 +117,7 @@ class StandardExecutor(Executor):
         listener = threading.Thread(target=log_listener, args=(log_queue,))
         listener.start()
 
-        with tqdm(**tqdm_kwargs) as pbar:
+        with self.pbar_class(**tqdm_kwargs) as pbar:
             pool = pool_class(
                 processes=max_workers,
                 initializer=initializer,
@@ -156,3 +156,8 @@ class StandardExecutor(Executor):
 @hookimpl
 def get_executor():
     return StandardExecutor()
+
+
+@hookimpl
+def get_progress_bar():
+    return tqdm

src/ocrmypdf/extra_plugins/awslambda.py

@@ -92,12 +92,10 @@ class LambdaExecutor(Executor):
         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)
+                task_finished(result, self.pbar_class)
             return
 
         self._lambda_pool_impl(
@@ -106,7 +104,6 @@ class LambdaExecutor(Executor):
             task=task,
             task_arguments=task_arguments,
             task_finished=task_finished,
-            pbar=pbar,
         )
 
     def _lambda_pool_impl(

src/ocrmypdf/pluginspec.py

@@ -65,14 +65,47 @@ def check_options(options: Namespace) -> None:
 
 @hookspec(firstresult=True)
 def get_executor() -> Executor:
-    """Called to perform parallel execution
+    """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.
+    ``conconcurrent.futures``, but they do not work the same way.
+
+    Should be of type :class:`Executor` or otherwise conforming to the protocol
+    of that call.
+
+    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_progress_bar():
+    """Called to obtain a class that can be used to create progress bars.
+
+    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.
+
+    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.
+
+    Here is how OCRmyPDF will use the progress bar:
+
+    Example:
+        pbar_class = pm.hook.get_progress_bar()
+        with pbar_class(**tqdm_kwargs) as pbar:
+            ...
+            pbar.update(1)
     """
 
 

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):