mirror of
https://github.com/Garmelon/PFERD.git
synced 2023-12-21 10:23:01 +01:00
Use conductor via context manager
This commit is contained in:
parent
6431a3fb3d
commit
2e85d26b6b
@ -1,6 +1,7 @@
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import AsyncIterator, Iterator, List, Optional
|
||||
from types import TracebackType
|
||||
from typing import AsyncIterator, Iterator, List, Optional, Type
|
||||
|
||||
import rich
|
||||
from rich.progress import Progress, TaskID
|
||||
@ -22,24 +23,30 @@ class TerminalConductor:
|
||||
self._progress = Progress()
|
||||
self._lines: List[str] = []
|
||||
|
||||
def _start(self) -> None:
|
||||
async def _start(self) -> None:
|
||||
async with self._lock:
|
||||
for line in self._lines:
|
||||
rich.print(line)
|
||||
self._lines = []
|
||||
|
||||
self._progress.start()
|
||||
|
||||
def _stop(self) -> None:
|
||||
async def _stop(self) -> None:
|
||||
async with self._lock:
|
||||
self._progress.stop()
|
||||
self._stopped = True
|
||||
|
||||
async def start(self) -> None:
|
||||
async with self._lock:
|
||||
self._start()
|
||||
async def __aenter__(self) -> None:
|
||||
await self._start()
|
||||
|
||||
async def stop(self) -> None:
|
||||
async with self._lock:
|
||||
self._stop()
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_value: Optional[BaseException],
|
||||
traceback: Optional[TracebackType],
|
||||
) -> Optional[bool]:
|
||||
await self._stop()
|
||||
return None
|
||||
|
||||
def print(self, line: str) -> None:
|
||||
if self._stopped:
|
||||
@ -50,11 +57,11 @@ class TerminalConductor:
|
||||
@asynccontextmanager
|
||||
async def exclusive_output(self) -> AsyncIterator[None]:
|
||||
async with self._lock:
|
||||
self.stop()
|
||||
self._stop()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.start()
|
||||
self._start()
|
||||
|
||||
@contextmanager
|
||||
def progress_bar(
|
||||
|
@ -63,11 +63,8 @@ class Crawler(ABC):
|
||||
return self.progress_bar(desc, total=size)
|
||||
|
||||
async def run(self) -> None:
|
||||
await self._conductor.start()
|
||||
try:
|
||||
async with self._conductor:
|
||||
await self.crawl()
|
||||
finally:
|
||||
await self._conductor.stop()
|
||||
|
||||
@abstractmethod
|
||||
async def crawl(self) -> None:
|
||||
|
Loading…
Reference in New Issue
Block a user