refactor: use a general progress factory

Author: NTFSvolume

src/mega/progress.py

@@ -2,70 +2,64 @@
 
 import asyncio
 import contextlib
-from collections.abc import Generator
 from contextvars import ContextVar
-from typing import TYPE_CHECKING, Literal, TypeAlias
-
-from rich.progress import BarColumn, DownloadColumn, Progress, SpinnerColumn, TimeRemainingColumn, TransferSpeedColumn
+from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
+    from types import TracebackType
 
-    ProgressHook: TypeAlias = Callable[[float], None]
+    from rich.progress import Progress
 
+    ProgressHook: TypeAlias = Callable[[float], None]
 
-_SHOW_PROGRESS = ContextVar[bool]("_SHOW_PROGRESS", default=False)
-_PROGRESS = ContextVar[Progress | None]("_PROGRESS", default=None)
+    class ProgressHookContext(Protocol):
+        def __enter__(self) -> ProgressHook: ...
 
+        def __exit__(
+            self, typ: type[BaseException] | None, value: BaseException | None, traceback: TracebackType | None, /
+        ) -> Any: ...
 
-def do_nothing(_: float) -> None: ...
+    class ProgressHookFactory(Protocol):
+        def __call__(self, description: str, total: float, kind: Literal["UP", "DOWN"]) -> ProgressHookContext: ...
 
 
-current_hook: ContextVar[ProgressHook] = ContextVar("current_hook", default=do_nothing)
+_PROGRESS_HOOK_FACTORY: ContextVar[ProgressHookFactory | None] = ContextVar("_PROGRESS_HOOK_FACTORY", default=None)
+current_hook: ContextVar[ProgressHook] = ContextVar("current_hook", default=lambda _: None)
 
 
 @contextlib.contextmanager
 def new_task(description: str, total: float, kind: Literal["UP", "DOWN"]) -> Generator[None]:
-    progress = _PROGRESS.get()
-    if progress is None:
+    factory = _PROGRESS_HOOK_FACTORY.get()
+    if factory is None:
         yield
         return
 
-    task_id = progress.add_task(description, total=total, kind=kind)
+    with factory(description, total, kind) as progress_hook:
+        token = current_hook.set(progress_hook)
+        try:
+            yield
+        finally:
+            current_hook.reset(token)
 
-    def progress_hook(advance: float) -> None:
-        progress.advance(task_id, advance)
 
-    token = current_hook.set(progress_hook)
-    try:
+@contextlib.contextmanager
+def new_progress() -> Generator[None]:
+    progress = _new_rich_progress()
+    if progress is None:
         yield
-    finally:
-        progress.remove_task(task_id=task_id)
-        current_hook.reset(token)
+        return
 
+    def hook_factory(*args, **kwargs):
+        return _new_rich_task(progress, *args, **kwargs)
+
+    token = _PROGRESS_HOOK_FACTORY.set(hook_factory)
 
-@contextlib.contextmanager
-def new_progress() -> Generator[None]:
-    progress = Progress(
-        "[{task.fields[kind]}]",
-        SpinnerColumn(),
-        "{task.description}",
-        BarColumn(bar_width=None),
-        "[progress.percentage]{task.percentage:>6.2f}%",
-        "-",
-        DownloadColumn(),
-        "-",
-        TransferSpeedColumn(),
-        "-",
-        TimeRemainingColumn(compact=True, elapsed_when_finished=True),
-        transient=True,
-    )
-    token = _PROGRESS.set(progress)
     try:
         with progress:
             yield
     finally:
-        _PROGRESS.reset(token)
+        _PROGRESS_HOOK_FACTORY.reset(token)
 
 
 async def test() -> None:
@@ -90,5 +84,50 @@ async def task(name: str) -> None:
                 tg.create_task(task(f"file{idx}"))
 
 
+def _new_rich_progress() -> Progress | None:
+    try:
+        from rich.progress import (
+            BarColumn,
+            DownloadColumn,
+            Progress,
+            SpinnerColumn,
+            TimeRemainingColumn,
+            TransferSpeedColumn,
+        )
+    except ImportError:
+        return None
+
+    else:
+        return Progress(
+            "[{task.fields[kind]}]",
+            SpinnerColumn(),
+            "{task.description}",
+            BarColumn(bar_width=None),
+            "[progress.percentage]{task.percentage:>6.2f}%",
+            "-",
+            DownloadColumn(),
+            "-",
+            TransferSpeedColumn(),
+            "-",
+            TimeRemainingColumn(compact=True, elapsed_when_finished=True),
+            transient=True,
+        )
+
+
+@contextlib.contextmanager
+def _new_rich_task(
+    progress: Progress, description: str, total: float, kind: Literal["UP", "DOWN"]
+) -> Generator[ProgressHook]:
+    task_id = progress.add_task(description, total=total, kind=kind)
+
+    def progress_hook(advance: float) -> None:
+        progress.advance(task_id, advance)
+
+    try:
+        yield progress_hook
+    finally:
+        progress.remove_task(task_id=task_id)
+
+
 if __name__ == "__main__":  # pragma: no coverage
     asyncio.run(test())

src/mega/utils.py

@@ -7,14 +7,14 @@
 from collections.abc import Awaitable, Iterable, Sequence
 from typing import Literal, TypeVar, overload
 
-from rich.logging import RichHandler
-
 from mega.errors import ValidationError
 
 _T = TypeVar("_T")
 
 
 def setup_logger(name: str = "mega") -> None:
+    from rich.logging import RichHandler
+
     handler = RichHandler(show_time=False, rich_tracebacks=True)
     logger = logging.getLogger(name)
     logger.setLevel(10)