Reorganize progress bars so they can be typed properly

src/ocrmypdf/_progressbar.py

@@ -0,0 +1,135 @@
+from typing import Protocol
+
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
+from rich.table import Column
+
+
+class ProgressBar(Protocol):
+    """The protocol that OCRmyPDF expects progress bar classes to be compatible with.
+
+    In practice this could be used for any time of monitoring, not just a progress bar.
+
+    Calling the class should return a new progress bar object, which is activated
+    with ``__enter__`` and terminated with ``__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.
+
+    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.
+    """
+
+    def __init__(
+        self,
+        *,
+        total: int | float | None,
+        desc: str | None,
+        unit: str | None,
+        disable: bool = False,
+        **kwargs,
+    ):
+        """Initialize a progress bar.
+
+        *total* indicates the total number of work units. If None, the total
+        number of work units is unknown. If *disable* is True, the progress bar
+        should be disabled. *unit* is a description of the work unit.
+        *desc* is a description of the overall task to be performed.
+
+        Unrecognized keyword arguments must be ignored, as the list of keyword
+        arguments may grow with time.
+        """
+
+    def __enter__(self):
+        """Enter a progress bar context."""
+
+    def __exit__(self, *args):
+        """Exit a progress bar context."""
+
+    def update(self, n=1):
+        """Update the progress bar by an increment.
+
+        For use within a progress bar context.
+        """
+
+
+class NullProgressBar:
+    """Progress bar API that takes no actions."""
+
+    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 RichProgressBar:
+    """Display progress bar using rich."""
+
+    def __init__(
+        self,
+        *,
+        console: Console,
+        desc: str,
+        total: float | None = None,
+        unit: str | None = None,
+        unit_scale: float | None = 1.0,
+        disable: bool = False,
+        **kwargs,
+    ):
+        self.progress = Progress(
+            TextColumn(
+                "[progress.description]{task.description}",
+                table_column=Column(min_width=20),
+            ),
+            BarColumn(),
+            TaskProgressColumn(),
+            MofNCompleteColumn(),
+            TimeRemainingColumn(),
+            console=console,
+            auto_refresh=True,
+            redirect_stderr=True,
+            redirect_stdout=False,
+            disable=disable,
+            **kwargs,
+        )
+        self.unit_scale = unit_scale
+        self.progress_bar = self.progress.add_task(
+            desc,
+            total=total * self.unit_scale
+            if total is not None and self.unit_scale is not None
+            else None,
+            unit=unit,
+        )
+
+    def __enter__(self):
+        self.progress.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.progress.refresh()
+        self.progress.stop()
+        return False
+
+    def update(self, value=None):
+        advance = self.unit_scale if value is None else value
+        self.progress.update(self.progress_bar, advance=advance)