42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import asyncio
|
|
import collections
|
|
from typing import Any, Deque, Optional
|
|
|
|
|
|
class EventResultOrError:
|
|
"""Event asyncio lock helper class.
|
|
|
|
Wraps the Event asyncio lock allowing either to awake the
|
|
locked Tasks without any error or raising an exception.
|
|
|
|
thanks to @vorpalsmith for the simple design.
|
|
"""
|
|
|
|
def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
|
|
self._loop = loop
|
|
self._exc: Optional[BaseException] = None
|
|
self._event = asyncio.Event()
|
|
self._waiters: Deque[asyncio.Future[Any]] = collections.deque()
|
|
|
|
def set(self, exc: Optional[BaseException] = None) -> None:
|
|
self._exc = exc
|
|
self._event.set()
|
|
|
|
async def wait(self) -> Any:
|
|
waiter = self._loop.create_task(self._event.wait())
|
|
self._waiters.append(waiter)
|
|
try:
|
|
val = await waiter
|
|
finally:
|
|
self._waiters.remove(waiter)
|
|
|
|
if self._exc is not None:
|
|
raise self._exc
|
|
|
|
return val
|
|
|
|
def cancel(self) -> None:
|
|
"""Cancel all waiters"""
|
|
for waiter in self._waiters:
|
|
waiter.cancel()
|