mirrors/pferd
1
0
Fork 0
mirror of https://github.com/Garmelon/PFERD.git synced 2023-12-21 10:23:01 +01:00
Code Issues Packages Projects Releases Wiki Activity

Clean up and fix conductor and limiter

Browse Source 
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
PFERD/conductor.py
View File

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

8
PFERD/limiter.py
View File

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