44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import asyncio
|
|
import collections
|
|
from typing import Any, Optional
|
|
|
|
try:
|
|
from typing import Deque
|
|
except ImportError:
|
|
from typing_extensions import Deque # noqa
|
|
|
|
|
|
class EventResultOrError:
|
|
"""
|
|
This class wrappers the Event asyncio lock allowing either 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 = None # type: Optional[BaseException]
|
|
self._event = asyncio.Event(loop=loop)
|
|
self._waiters = collections.deque() # type: Deque[asyncio.Future[Any]]
|
|
|
|
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()
|