wehub-resource-sync
129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
import asyncio
|
|
from collections import deque
|
|
from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterator
|
|
from typing import Any, Generic, Protocol, TypeVar, overload, runtime_checkable
|
|
|
|
from typing_extensions import AsyncContextManager
|
|
|
|
# based on https://github.com/maxfischer2781/asyncstdlib/blob/master/asyncstdlib/itertools.py
|
|
|
|
|
|
@runtime_checkable
|
|
class _ACloseable(Protocol):
|
|
async def aclose(self) -> None:
|
|
"""Asynchronously close this object"""
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
async def tee_peer(
|
|
iterator: AsyncIterator[T],
|
|
buffer: deque[T],
|
|
peers: list[deque[T]],
|
|
lock: AsyncContextManager[Any],
|
|
exception: list[BaseException | None],
|
|
) -> AsyncGenerator[T, None]:
|
|
# exception is a shared mutable container across all peers. When the upstream
|
|
# iterator raises, only the first peer to call __anext__() would normally see
|
|
# the error — subsequent calls return StopAsyncIteration per Python async
|
|
# generator semantics, silently swallowing the error for other peers.
|
|
#
|
|
# To fix this, the first peer to hit the exception stores it in exception[0].
|
|
# Other peers check this before advancing the iterator and re-raise the same
|
|
# exception, ensuring all peers observe the upstream failure.
|
|
try:
|
|
while True:
|
|
if not buffer:
|
|
async with lock:
|
|
if buffer:
|
|
continue
|
|
# a peer already hit an upstream error — re-raise for this peer
|
|
if exception[0] is not None:
|
|
raise exception[0]
|
|
try:
|
|
item = await iterator.__anext__()
|
|
except StopAsyncIteration:
|
|
break
|
|
except asyncio.CancelledError:
|
|
# CancelledError is task-specific — don't store it in the
|
|
# shared exception list to avoid cascading to other peers
|
|
raise
|
|
except BaseException as e:
|
|
exception[0] = e
|
|
raise
|
|
else:
|
|
for peer_buffer in peers:
|
|
peer_buffer.append(item)
|
|
yield buffer.popleft()
|
|
finally:
|
|
for idx, peer_buffer in enumerate(peers): # pragma: no branch
|
|
if peer_buffer is buffer:
|
|
peers.pop(idx)
|
|
break
|
|
|
|
if not peers and isinstance(iterator, _ACloseable):
|
|
await iterator.aclose()
|
|
|
|
|
|
class Tee(Generic[T]):
|
|
__slots__ = ("_iterator", "_buffers", "_children")
|
|
|
|
def __init__(
|
|
self,
|
|
iterator: AsyncIterable[T],
|
|
n: int = 2,
|
|
):
|
|
self._iterator = iterator.__aiter__()
|
|
self._buffers: list[deque[T]] = [deque() for _ in range(n)]
|
|
|
|
lock = asyncio.Lock()
|
|
exception: list[BaseException | None] = [None]
|
|
self._children = tuple(
|
|
tee_peer(
|
|
iterator=self._iterator,
|
|
buffer=buffer,
|
|
peers=self._buffers,
|
|
lock=lock,
|
|
exception=exception,
|
|
)
|
|
for buffer in self._buffers
|
|
)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._children)
|
|
|
|
@overload
|
|
def __getitem__(self, item: int) -> AsyncIterator[T]: ...
|
|
|
|
@overload
|
|
def __getitem__(self, item: slice) -> tuple[AsyncIterator[T], ...]: ...
|
|
|
|
def __getitem__(self, item: int | slice) -> AsyncIterator[T] | tuple[AsyncIterator[T], ...]:
|
|
return self._children[item]
|
|
|
|
def __iter__(self) -> Iterator[AsyncIterator[T]]:
|
|
yield from self._children
|
|
|
|
async def __aenter__(self) -> "Tee[T]":
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
await self.aclose()
|
|
|
|
async def aclose(self) -> None:
|
|
for child in self._children:
|
|
try:
|
|
await child.aclose()
|
|
except Exception:
|
|
pass
|
|
|
|
if isinstance(self._iterator, _ACloseable):
|
|
try:
|
|
await self._iterator.aclose()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
tee = Tee
|