From 19045c4f218ed2dc7ab264fb724b55dcf58a9469 Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <james@purplerock.ca>
Date: Sat, 3 Jun 2023 00:29:26 -0700
Subject: [PATCH] Replace coloredlogs and tqdm with rich

---
 pyproject.toml                              |  2 +-
 src/ocrmypdf/_logging.py                    | 63 +++++++++++++++++++++
 src/ocrmypdf/api.py                         | 10 ----
 src/ocrmypdf/builtin_plugins/concurrency.py | 15 +++--
 4 files changed, 75 insertions(+), 15 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 5184ad86..021e33ba 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,7 +17,6 @@ license = {text = "MPL-2.0"}
 requires-python = ">=3.8"
 dependencies = [
   "Pillow>=8.2.0",
-  "coloredlogs>=14.0",
   "deprecation>=2.1.0",
   "img2pdf>=0.3.0", # pure Python
   "packaging>=20",
@@ -25,6 +24,7 @@ dependencies = [
   "pikepdf>=5.0.1",
   "pluggy>=0.13.0",
   "reportlab>=3.5.66",
+  "rich>=13",
   "tqdm>=4",
   "importlib-resources>=5;python_version<'3.9'",  # until Python 3.9
   "typing-extensions>=4;python_version<'3.10'",
diff --git a/src/ocrmypdf/_logging.py b/src/ocrmypdf/_logging.py
index ef939d71..d59685c6 100644
--- a/src/ocrmypdf/_logging.py
+++ b/src/ocrmypdf/_logging.py
@@ -8,6 +8,16 @@ from __future__ import annotations
 import logging
 from contextlib import suppress
 
+from rich.console import Console
+from rich.logging import RichHandler
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
 from tqdm import tqdm
 
 
@@ -41,3 +51,56 @@ class TqdmConsole:
     def flush(self):
         with suppress(AttributeError):
             self.file.flush()
+
+
+class RichLoggingHandler(RichHandler):
+    def __init__(self, console: Console, **kwargs):
+        super().__init__(
+            console=console, show_level=False, show_time=False, markup=True, **kwargs
+        )
+
+
+class RichTqdmProgressAdapter:
+    """Adapt tqdm API to rich progress bar."""
+
+    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}"),
+            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, 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)
diff --git a/src/ocrmypdf/api.py b/src/ocrmypdf/api.py
index 1c6d01eb..b0f8c455 100644
--- a/src/ocrmypdf/api.py
+++ b/src/ocrmypdf/api.py
@@ -15,9 +15,6 @@ from pathlib import Path
 from typing import AnyStr, BinaryIO, Iterable, Union
 from warnings import warn
 
-import coloredlogs
-from humanfriendly.terminal import enable_ansi_support
-
 from ocrmypdf._logging import PageNumberFilter, TqdmConsole
 from ocrmypdf._plugin_manager import get_plugin_manager
 from ocrmypdf._sync import run_pipeline
@@ -112,14 +109,7 @@ def configure_logging(
     else:
         fmt = '%(pageno)s%(message)s'
 
-    use_colors = progress_bar_friendly
     formatter = None
-    if use_colors:
-        use_colors = enable_ansi_support()
-        if use_colors:
-            use_colors = coloredlogs.terminal_supports_colors()
-        if use_colors:
-            formatter = coloredlogs.ColoredFormatter(fmt=fmt)
 
     if not formatter:
         formatter = logging.Formatter(fmt=fmt)
diff --git a/src/ocrmypdf/builtin_plugins/concurrency.py b/src/ocrmypdf/builtin_plugins/concurrency.py
index d9b186f8..9c08b526 100644
--- a/src/ocrmypdf/builtin_plugins/concurrency.py
+++ b/src/ocrmypdf/builtin_plugins/concurrency.py
@@ -16,10 +16,10 @@ from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_compl
 from contextlib import suppress
 from typing import Callable, Iterable, Type, Union
 
-from tqdm import tqdm
+from rich.console import Console as RichConsole
 
 from ocrmypdf import Executor, hookimpl
-from ocrmypdf._logging import TqdmConsole
+from ocrmypdf._logging import RichLoggingHandler, RichTqdmProgressAdapter
 from ocrmypdf.exceptions import InputFileError
 from ocrmypdf.helpers import remove_all_log_handlers
 
@@ -168,13 +168,20 @@ def get_executor(progressbar_class):
     return StandardExecutor(pbar_class=progressbar_class)
 
 
+RICH_CONSOLE = RichConsole(stderr=True)
+
+
 @hookimpl
 def get_progressbar_class():
     """Return the default progress bar class."""
-    return tqdm
+
+    def partial_RichTqdmProgressAdapter(*args, **kwargs):
+        return RichTqdmProgressAdapter(*args, **kwargs, console=RICH_CONSOLE)
+
+    return partial_RichTqdmProgressAdapter
 
 
 @hookimpl
 def get_logging_console():
     """Return the default logging console handler."""
-    return logging.StreamHandler(stream=TqdmConsole(sys.stderr))
+    return RichLoggingHandler(console=RICH_CONSOLE)