pferd/PFERD/conductor.py

76 lines
1.9 KiB
Python
Raw Normal View History

2021-04-29 11:25:00 +02:00
import asyncio
from contextlib import asynccontextmanager, contextmanager
from typing import AsyncIterator, Iterator, List, Optional
2021-04-29 11:25:00 +02:00
import rich
from rich.progress import Progress, TaskID
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)
class TerminalConductor:
def __init__(self) -> None:
self._stopped = False
self._lock = asyncio.Lock()
self._progress = Progress()
self._lines: List[str] = []
def _start(self) -> None:
for line in self._lines:
rich.print(line)
self._lines = []
self._progress.start()
def _stop(self) -> None:
self._progress.stop()
self._stopped = True
async def start(self) -> None:
async with self._lock:
2021-04-29 11:25:00 +02:00
self._start()
async def stop(self) -> None:
async with self._lock:
2021-04-29 11:25:00 +02:00
self._stop()
def print(self, line: str) -> None:
if self._stopped:
self._lines.append(line)
else:
rich.print(line)
@asynccontextmanager
async def exclusive_output(self) -> AsyncIterator[None]:
2021-04-29 11:25:00 +02:00
async with self._lock:
self.stop()
try:
yield
finally:
self.start()
@contextmanager
def progress_bar(
2021-04-29 11:25:00 +02:00
self,
description: str,
total: Optional[float] = None,
2021-04-29 11:25:00 +02:00
) -> Iterator[ProgressBar]:
if total is None:
# Indeterminate progress bar
taskid = self._progress.add_task(description, start=False)
else:
taskid = self._progress.add_task(description, total=total)
2021-04-29 11:25:00 +02:00
bar = ProgressBar(self._progress, taskid)
try:
yield bar
finally:
self._progress.remove_task(taskid)