codekingpro/portable-devtools

1"""Asynchronous iterator utilities.2 3Adapted from4https://github.com/maxfischer2781/asyncstdlib/blob/master/asyncstdlib/itertools.py5MIT License.6"""7 8from collections import deque9from collections.abc import (10    AsyncGenerator,11    AsyncIterable,12    AsyncIterator,13    Awaitable,14    Callable,15    Iterator,16)17from contextlib import AbstractAsyncContextManager18from types import TracebackType19from typing import (20    Any,21    Generic,22    TypeVar,23    cast,24    overload,25)26 27from typing_extensions import override28 29from langchain_core._api.deprecation import deprecated30 31T = TypeVar("T")32 33_no_default = object()34 35 36# https://github.com/python/cpython/blob/main/Lib/test/test_asyncgen.py#L5437@deprecated(since="1.1.2", removal="2.0.0")38def py_anext(39    iterator: AsyncIterator[T], default: T | Any = _no_default40) -> Awaitable[T | Any | None]:41    """Pure-Python implementation of `anext()` for testing purposes.42 43    Closely matches the builtin `anext()` C implementation.44 45    Can be used to compare the built-in implementation of the inner coroutines machinery46    to C-implementation of `__anext__()` and `send()` or `throw()` on the returned47    generator.48 49    Args:50        iterator: The async iterator to advance.51        default: The value to return if the iterator is exhausted.52 53            If not provided, a `StopAsyncIteration` exception is raised.54 55    Returns:56        The next value from the iterator, or the default value if the iterator is57            exhausted.58 59    Raises:60        TypeError: If the iterator is not an async iterator.61    """62    try:63        __anext__ = cast(64            "Callable[[AsyncIterator[T]], Awaitable[T]]", type(iterator).__anext__65        )66    except AttributeError as e:67        msg = f"{iterator!r} is not an async iterator"68        raise TypeError(msg) from e69 70    if default is _no_default:71        return __anext__(iterator)72 73    async def anext_impl() -> T | Any:74        try:75            # The C code is way more low-level than this, as it implements76            # all methods of the iterator protocol. In this implementation77            # we're relying on higher-level coroutine concepts, but that's78            # exactly what we want -- crosstest pure-Python high-level79            # implementation and low-level C anext() iterators.80            return await __anext__(iterator)81        except StopAsyncIteration:82            return default83 84    return anext_impl()85 86 87class NoLock:88    """Dummy lock that provides the proper interface but no protection."""89 90    async def __aenter__(self) -> None:91        """Do nothing."""92 93    async def __aexit__(94        self,95        exc_type: type[BaseException] | None,96        exc_val: BaseException | None,97        exc_tb: TracebackType | None,98    ) -> bool:99        """Return False, exception not suppressed."""100        return False101 102 103async def tee_peer(104    iterator: AsyncIterator[T],105    # the buffer specific to this peer106    buffer: deque[T],107    # the buffers of all peers, including our own108    peers: list[deque[T]],109    lock: AbstractAsyncContextManager[Any],110) -> AsyncGenerator[T, None]:111    """An individual iterator of a `tee`.112 113    This function is a generator that yields items from the shared iterator114    `iterator`. It buffers items until the least advanced iterator has yielded them as115    well.116 117    The buffer is shared with all other peers.118 119    Args:120        iterator: The shared iterator.121        buffer: The buffer for this peer.122        peers: The buffers of all peers.123        lock: The lock to synchronise access to the shared buffers.124 125    Yields:126        The next item from the shared iterator.127    """128    try:129        while True:130            if not buffer:131                async with lock:132                    # Another peer produced an item while we were waiting for the lock.133                    # Proceed with the next loop iteration to yield the item.134                    if buffer:135                        continue136                    try:137                        item = await anext(iterator)138                    except StopAsyncIteration:139                        break140                    else:141                        # Append to all buffers, including our own. We'll fetch our142                        # item from the buffer again, instead of yielding it directly.143                        # This ensures the proper item ordering if any of our peers144                        # are fetching items concurrently. They may have buffered their145                        # item already.146                        for peer_buffer in peers:147                            peer_buffer.append(item)148            yield buffer.popleft()149    finally:150        async with lock:151            # this peer is done - remove its buffer152            for idx, peer_buffer in enumerate(peers):  # pragma: no branch153                if peer_buffer is buffer:154                    peers.pop(idx)155                    break156            # if we are the last peer, try and close the iterator157            if not peers and hasattr(iterator, "aclose"):158                await iterator.aclose()159 160 161class Tee(Generic[T]):162    """Create `n` separate asynchronous iterators over `iterable`.163 164    This splits a single `iterable` into multiple iterators, each providing165    the same items in the same order.166 167    All child iterators may advance separately but share the same items from `iterable`168    -- when the most advanced iterator retrieves an item, it is buffered until the least169    advanced iterator has yielded it as well.170 171    A `tee` works lazily and can handle an infinite `iterable`, provided172    that all iterators advance.173 174    ```python175    async def derivative(sensor_data):176        previous, current = a.tee(sensor_data, n=2)177        await a.anext(previous)  # advance one iterator178        return a.map(operator.sub, previous, current)179    ```180 181    Unlike `itertools.tee`, `.tee` returns a custom type instead of a `tuple`. Like a182    tuple, it can be indexed, iterated and unpacked to get the child iterators. In183    addition, its `.tee.aclose` method immediately closes all children, and it can be184    used in an `async with` context for the same effect.185 186    If `iterable` is an iterator and read elsewhere, `tee` will *not* provide these187    items. Also, `tee` must internally buffer each item until the last iterator has188    yielded it; if the most and least advanced iterator differ by most data, using a189    `list` is more efficient (but not lazy).190 191    If the underlying iterable is concurrency safe (`anext` may be awaited concurrently)192    the resulting iterators are concurrency safe as well. Otherwise, the iterators are193    safe if there is only ever one single "most advanced" iterator.194 195    To enforce sequential use of `anext`, provide a `lock`196 197    - e.g. an `asyncio.Lock` instance in an `asyncio` application - and access is198        automatically synchronised.199 200    """201 202    def __init__(203        self,204        iterable: AsyncIterator[T],205        n: int = 2,206        *,207        lock: AbstractAsyncContextManager[Any] | None = None,208    ):209        """Create a `tee`.210 211        Args:212            iterable: The iterable to split.213            n: The number of iterators to create.214            lock: The lock to synchronise access to the shared buffers.215 216        """217        self._iterator = iterable.__aiter__()  # before 3.10 aiter() doesn't exist218        self._buffers: list[deque[T]] = [deque() for _ in range(n)]219        self._children = tuple(220            tee_peer(221                iterator=self._iterator,222                buffer=buffer,223                peers=self._buffers,224                lock=lock if lock is not None else NoLock(),225            )226            for buffer in self._buffers227        )228 229    def __len__(self) -> int:230        """Return the number of child iterators."""231        return len(self._children)232 233    @overload234    def __getitem__(self, item: int) -> AsyncIterator[T]: ...235 236    @overload237    def __getitem__(self, item: slice) -> tuple[AsyncIterator[T], ...]: ...238 239    def __getitem__(240        self, item: int | slice241    ) -> AsyncIterator[T] | tuple[AsyncIterator[T], ...]:242        """Return the child iterator(s) for the given index or slice."""243        return self._children[item]244 245    def __iter__(self) -> Iterator[AsyncIterator[T]]:246        """Iterate over the child iterators.247 248        Yields:249            The child iterators.250        """251        yield from self._children252 253    async def __aenter__(self) -> "Tee[T]":254        """Return the tee instance."""255        return self256 257    async def __aexit__(258        self,259        exc_type: type[BaseException] | None,260        exc_val: BaseException | None,261        exc_tb: TracebackType | None,262    ) -> bool:263        """Close all child iterators.264 265        Returns:266            `False`, exceptions not suppressed.267        """268        await self.aclose()269        return False270 271    async def aclose(self) -> None:272        """Async close all child iterators."""273        for child in self._children:274            await child.aclose()275 276 277atee = Tee278 279 280class aclosing(AbstractAsyncContextManager):  # noqa: N801281    """Async context manager to wrap an `AsyncGenerator` that has a `aclose()` method.282 283    Code like this:284 285    ```python286    async with aclosing(<module>.fetch(<arguments>)) as agen:287        <block>288    ```289 290    ...is equivalent to this:291 292    ```python293    agen = <module>.fetch(<arguments>)294    try:295        <block>296    finally:297        await agen.aclose()298 299    ```300    """301 302    def __init__(self, thing: AsyncGenerator[Any, Any] | AsyncIterator[Any]) -> None:303        """Create the context manager.304 305        Args:306            thing: The resource to wrap.307        """308        self.thing = thing309 310    @override311    async def __aenter__(self) -> AsyncGenerator[Any, Any] | AsyncIterator[Any]:312        return self.thing313 314    @override315    async def __aexit__(316        self,317        exc_type: type[BaseException] | None,318        exc_value: BaseException | None,319        traceback: TracebackType | None,320    ) -> None:321        if hasattr(self.thing, "aclose"):322            await self.thing.aclose()323 324 325async def abatch_iterate(326    size: int, iterable: AsyncIterable[T]327) -> AsyncIterator[list[T]]:328    """Utility batching function for async iterables.329 330    Args:331        size: The size of the batch.332        iterable: The async iterable to batch.333 334    Yields:335        The batches.336    """337    batch: list[T] = []338    async for element in iterable:339        if len(batch) < size:340            batch.append(element)341 342        if len(batch) >= size:343            yield batch344            batch = []345 346    if batch:347        yield batch348