import asyncio from contextlib import asynccontextmanager, contextmanager # TODO In Python 3.9 and above, ContextManager and AsyncContextManager are deprecated from typing import AsyncIterator, ContextManager, Iterator, List, Optional from rich.console import Console, RenderGroup from rich.live import Live from rich.markup import escape from rich.progress import (BarColumn, DownloadColumn, Progress, TaskID, TextColumn, TimeRemainingColumn, TransferSpeedColumn) from rich.table import Column class ProgressBar: def __init__(self, progress: Progress, taskid: TaskID): self._progress = progress self._taskid = taskid def advance(self, amount: float = 1) -> None: self._progress.advance(self._taskid, advance=amount) def set_total(self, total: float) -> None: self._progress.update(self._taskid, total=total) self._progress.start_task(self._taskid) class Log: def __init__(self) -> None: self.console = Console(highlight=False) self._crawl_progress = Progress( TextColumn("{task.description}", table_column=Column(ratio=1)), BarColumn(), TimeRemainingColumn(), expand=True, ) self._download_progress = Progress( TextColumn("{task.description}", table_column=Column(ratio=1)), TransferSpeedColumn(), DownloadColumn(), BarColumn(), TimeRemainingColumn(), expand=True, ) self._live = Live(console=self.console, transient=True) self._update_live() self._showing_progress = False self._progress_suspended = False self._lock = asyncio.Lock() self._lines: List[str] = [] # Whether different parts of the output are enabled or disabled self.output_explain = False self.output_action = True self.output_report = True def _update_live(self) -> None: elements = [] if self._crawl_progress.task_ids: elements.append(self._crawl_progress) if self._download_progress.task_ids: elements.append(self._download_progress) group = RenderGroup(*elements) # type: ignore self._live.update(group) @contextmanager def show_progress(self) -> Iterator[None]: if self._showing_progress: raise RuntimeError("Calling 'show_progress' while already showing progress") self._showing_progress = True try: with self._live: yield finally: self._showing_progress = False @asynccontextmanager async def exclusive_output(self) -> AsyncIterator[None]: if not self._showing_progress: raise RuntimeError("Calling 'exclusive_output' while not showing progress") async with self._lock: self._progress_suspended = True self._live.stop() try: yield finally: self._live.start() self._progress_suspended = False for line in self._lines: self.print(line) self._lines = [] def print(self, text: str) -> None: if self._progress_suspended: self._lines.append(text) else: self.console.print(text) def explain_topic(self, text: str) -> None: if self.output_explain: self.print(f"[cyan]{escape(text)}") def explain(self, text: str) -> None: if self.output_explain: self.print(f" {escape(text)}") def action(self, text: str) -> None: if self.output_action: self.print(text) def report(self, text: str) -> None: if self.output_report: self.print(text) @contextmanager def _bar( self, progress: Progress, description: str, total: Optional[float], ) -> Iterator[ProgressBar]: if total is None: # Indeterminate progress bar taskid = progress.add_task(description, start=False) else: taskid = progress.add_task(description, total=total) self._update_live() try: yield ProgressBar(progress, taskid) finally: progress.remove_task(taskid) self._update_live() def crawl_bar( self, description: str, total: Optional[float] = None, ) -> ContextManager[ProgressBar]: return self._bar(self._crawl_progress, description, total) def download_bar( self, description: str, total: Optional[float] = None, ) -> ContextManager[ProgressBar]: return self._bar(self._download_progress, description, total) log = Log()