wehub-resource-sync
179 lines
4.6 KiB
Python
179 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
from collections import deque
|
|
from collections.abc import AsyncIterator
|
|
from typing import Generic, Protocol, TypeVar
|
|
|
|
T = TypeVar("T")
|
|
T_co = TypeVar("T_co", covariant=True)
|
|
T_contra = TypeVar("T_contra", contravariant=True)
|
|
|
|
|
|
# Based on asyncio.Queue, see https://github.com/python/cpython/blob/main/Lib/asyncio/queues.py
|
|
|
|
|
|
class ChanClosed(Exception):
|
|
pass
|
|
|
|
|
|
class ChanFull(Exception):
|
|
pass
|
|
|
|
|
|
class ChanEmpty(Exception):
|
|
pass
|
|
|
|
|
|
class ChanSender(Protocol[T_contra]):
|
|
async def send(self, value: T_contra) -> None: ...
|
|
|
|
def send_nowait(self, value: T_contra) -> None: ...
|
|
|
|
def close(self) -> None: ...
|
|
|
|
|
|
class ChanReceiver(Protocol[T_co]):
|
|
async def recv(self) -> T_co: ...
|
|
|
|
def recv_nowait(self) -> T_co: ...
|
|
|
|
def close(self) -> None: ...
|
|
|
|
def __aiter__(self) -> AsyncIterator[T_co]: ...
|
|
|
|
async def __anext__(self) -> T_co: ...
|
|
|
|
|
|
class Chan(Generic[T]):
|
|
def __init__(
|
|
self,
|
|
maxsize: int = 0,
|
|
loop: asyncio.AbstractEventLoop | None = None,
|
|
) -> None:
|
|
self._loop = loop or asyncio.get_event_loop()
|
|
self._maxsize = max(maxsize, 0)
|
|
# self._finished_ev = asyncio.Event()
|
|
self._close_ev = asyncio.Event()
|
|
self._closed = False
|
|
self._gets: deque[asyncio.Future[T | None]] = deque()
|
|
self._puts: deque[asyncio.Future[T | None]] = deque()
|
|
self._queue: deque[T] = deque()
|
|
|
|
def _wakeup_next(self, waiters: deque[asyncio.Future[T | None]]) -> None:
|
|
while waiters:
|
|
waiter = waiters.popleft()
|
|
if not waiter.done():
|
|
waiter.set_result(None)
|
|
break
|
|
|
|
async def send(self, value: T) -> None:
|
|
while self.full() and not self._close_ev.is_set():
|
|
p = self._loop.create_future()
|
|
self._puts.append(p)
|
|
try:
|
|
await p
|
|
except ChanClosed:
|
|
raise
|
|
except:
|
|
p.cancel()
|
|
with contextlib.suppress(ValueError):
|
|
self._puts.remove(p)
|
|
|
|
if not self.full() and not p.cancelled():
|
|
self._wakeup_next(self._puts)
|
|
raise
|
|
|
|
self.send_nowait(value)
|
|
|
|
def send_nowait(self, value: T) -> None:
|
|
if self._close_ev.is_set():
|
|
raise ChanClosed
|
|
|
|
if self.full():
|
|
raise ChanFull
|
|
|
|
self._queue.append(value)
|
|
self._wakeup_next(self._gets)
|
|
|
|
async def recv(self) -> T:
|
|
while self.empty() and not self._close_ev.is_set():
|
|
g = self._loop.create_future()
|
|
self._gets.append(g)
|
|
|
|
try:
|
|
await g
|
|
except ChanClosed:
|
|
raise
|
|
except BaseException:
|
|
g.cancel()
|
|
with contextlib.suppress(ValueError):
|
|
self._gets.remove(g)
|
|
|
|
if not self.empty() and not g.cancelled():
|
|
self._wakeup_next(self._gets)
|
|
|
|
raise
|
|
|
|
return self.recv_nowait()
|
|
|
|
def recv_nowait(self) -> T:
|
|
if self.empty():
|
|
if self._close_ev.is_set():
|
|
raise ChanClosed
|
|
else:
|
|
raise ChanEmpty
|
|
item = self._queue.popleft()
|
|
# if self.empty() and self._close_ev.is_set():
|
|
# self._finished_ev.set()
|
|
self._wakeup_next(self._puts)
|
|
return item
|
|
|
|
def close(self) -> None:
|
|
self._closed = True
|
|
self._close_ev.set()
|
|
for putter in self._puts:
|
|
if not putter.cancelled():
|
|
putter.set_exception(ChanClosed())
|
|
|
|
while len(self._gets) > self.qsize():
|
|
getter = self._gets.pop()
|
|
if not getter.cancelled():
|
|
getter.set_exception(ChanClosed())
|
|
|
|
while self._gets:
|
|
self._wakeup_next(self._gets)
|
|
|
|
# if self.empty():
|
|
# self._finished_ev.set()
|
|
|
|
@property
|
|
def closed(self) -> bool:
|
|
return self._closed
|
|
|
|
# async def join(self) -> None:
|
|
# await self._finished_ev.wait()
|
|
|
|
def qsize(self) -> int:
|
|
"""the number of elements queued (unread) in the channel buffer"""
|
|
return len(self._queue)
|
|
|
|
def full(self) -> bool:
|
|
if self._maxsize <= 0:
|
|
return False
|
|
else:
|
|
return self.qsize() >= self._maxsize
|
|
|
|
def empty(self) -> bool:
|
|
return not self._queue
|
|
|
|
def __aiter__(self) -> AsyncIterator[T]:
|
|
return self
|
|
|
|
async def __anext__(self) -> T:
|
|
try:
|
|
return await self.recv()
|
|
except ChanClosed:
|
|
raise StopAsyncIteration from None
|