46 lines
1.8 KiB
Python
46 lines
1.8 KiB
Python
"""Terminal progress on stderr; dataset statistics remain machine-readable on stdout."""
|
||
from rich.console import Console
|
||
from rich.progress import (
|
||
BarColumn, MofNCompleteColumn, Progress, ProgressColumn,
|
||
TaskProgressColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn,
|
||
)
|
||
from rich.text import Text
|
||
|
||
|
||
class RateColumn(ProgressColumn):
|
||
def render(self, task):
|
||
speed = task.finished_speed or task.speed
|
||
return Text(f'{speed:.1f} зап/с' if speed is not None else '— зап/с')
|
||
|
||
|
||
class ProgressDisplay:
|
||
def __init__(self, label, total, *, initial=0, enabled=True):
|
||
self.label = label
|
||
self.total = total
|
||
self.progress = Progress(
|
||
TextColumn('{task.description}'), BarColumn(bar_width=None),
|
||
TaskProgressColumn(), MofNCompleteColumn(), RateColumn(),
|
||
TextColumn('прошло'), TimeElapsedColumn(),
|
||
TextColumn('осталось'), TimeRemainingColumn(),
|
||
console=Console(stderr=True), refresh_per_second=4,
|
||
redirect_stdout=False, redirect_stderr=False, disable=not enabled,
|
||
)
|
||
self.task = self.progress.add_task(label, total=total, completed=min(initial, total))
|
||
|
||
def __enter__(self):
|
||
self.progress.start()
|
||
return self
|
||
|
||
def update(self, completed=None, *, status=None):
|
||
kwargs = {}
|
||
if completed is not None:
|
||
kwargs['completed'] = min(completed, self.total)
|
||
if status is not None:
|
||
kwargs['description'] = f'{self.label}: {status}' if status else self.label
|
||
self.progress.update(self.task, **kwargs)
|
||
|
||
def __exit__(self, exc_type, exc, traceback):
|
||
if exc_type is not None:
|
||
self.update(status='прервано' if issubclass(exc_type, (KeyboardInterrupt, SystemExit)) else 'ошибка')
|
||
self.progress.stop()
|