Clean up and fix conductor and limiter

Turns out you have to await an async lock, who knew...
This commit is contained in:
Joscha 2021-04-29 13:43:50 +02:00
parent c4fb92c658
commit 7e127cd5cc
2 changed files with 14 additions and 27 deletions

View File

@ -1,12 +1,8 @@
import asyncio
from contextlib import asynccontextmanager, contextmanager
from pathlib import Path
# TODO If we upgrade to python 3.9, these context manager hints are deprecated
from typing import (AsyncContextManager, AsyncIterator, ContextManager,
Iterator, List, Optional)
from typing import AsyncIterator, Iterator, List, Optional
import rich
from rich.markup import escape
from rich.progress import Progress, TaskID
@ -38,11 +34,11 @@ class TerminalConductor:
self._stopped = True
async def start(self) -> None:
with self._lock:
async with self._lock:
self._start()
async def stop(self) -> None:
with self._lock:
async with self._lock:
self._stop()
def print(self, line: str) -> None:
@ -52,7 +48,7 @@ class TerminalConductor:
rich.print(line)
@asynccontextmanager
async def _exclusive_output_cm(self) -> AsyncIterator[None]:
async def exclusive_output(self) -> AsyncIterator[None]:
async with self._lock:
self.stop()
try:
@ -60,25 +56,20 @@ class TerminalConductor:
finally:
self.start()
def exclusive_output(self) -> AsyncContextManager[None]:
return self._exclusive_output_cm()
@contextmanager
def _progress_bar_cm(
def progress_bar(
self,
description: str,
steps: Optional[float],
total: Optional[float] = None,
) -> Iterator[ProgressBar]:
taskid = self._progress.add_task(description, steps=steps)
if total is None:
# Indeterminate progress bar
taskid = self._progress.add_task(description, start=False)
else:
taskid = self._progress.add_task(description, total=total)
bar = ProgressBar(self._progress, taskid)
try:
yield bar
finally:
self._progress.remove_task(taskid)
def progress_bar(
self,
description: Path,
steps: Optional[float],
) -> ContextManager[ProgressBar]:
return self._progress_bar_cm(escape(str(description)), steps=steps)

View File

@ -1,7 +1,6 @@
import asyncio
from contextlib import asynccontextmanager
# TODO If we upgrade to python 3.9, this context manager hint is deprecated
from typing import AsyncContextManager, AsyncIterator
from typing import AsyncIterator
class Limiter:
@ -9,12 +8,9 @@ class Limiter:
self._semaphore = asyncio.Semaphore(limit)
@asynccontextmanager
async def _context_manager(self) -> AsyncIterator[None]:
async def limit(self) -> AsyncIterator[None]:
await self._semaphore.acquire()
try:
yield
finally:
self._semaphore.release()
def limit(self) -> AsyncContextManager[None]:
return self._context_manager()