обновлен формат генерации датасета

assembly_dataset/progress.py

@@ -0,0 +1,45 @@
+"""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()