codekingpro/portable-devtools

1"""Adapted.2 3Original source:4https://github.com/maxfischer2781/asyncstdlib/blob/master/asyncstdlib/itertools.py5MIT License6"""7 8from __future__ import annotations9 10import asyncio11import contextvars12import functools13import inspect14from collections import deque15from collections.abc import (16    AsyncGenerator,17    AsyncIterable,18    AsyncIterator,19    Awaitable,20    Coroutine,21    Iterable,22    Iterator,23)24from contextlib import AbstractAsyncContextManager25from typing import (26    Any,27    Callable,28    Generic,29    Optional,30    TypeVar,31    Union,32    cast,33    overload,34)35 36from langsmith._runtime_overrides import get_runtime_overrides37 38T = TypeVar("T")39 40_no_default = object()41 42 43# https://github.com/python/cpython/blob/main/Lib/test/test_asyncgen.py#L5444# before 3.10, the builtin anext() was not available45def py_anext(46    iterator: AsyncIterator[T], default: Union[T, Any] = _no_default47) -> Awaitable[Union[T, None, Any]]:48    """Pure-Python implementation of anext() for testing purposes.49 50    Closely matches the builtin anext() C implementation.51    Can be used to compare the built-in implementation of the inner52    coroutines machinery to C-implementation of __anext__() and send()53    or throw() on the returned generator.54    """55    try:56        __anext__ = cast(57            Callable[[AsyncIterator[T]], Awaitable[T]], type(iterator).__anext__58        )59    except AttributeError:60        raise TypeError(f"{iterator!r} is not an async iterator")61 62    if default is _no_default:63        return __anext__(iterator)64 65    async def anext_impl() -> Union[T, Any]:66        try:67            # The C code is way more low-level than this, as it implements68            # all methods of the iterator protocol. In this implementation69            # we're relying on higher-level coroutine concepts, but that's70            # exactly what we want -- crosstest pure-Python high-level71            # implementation and low-level C anext() iterators.72            return await __anext__(iterator)73        except StopAsyncIteration:74            return default75 76    return anext_impl()77 78 79class NoLock:80    """Dummy lock that provides the proper interface but no protection."""81 82    async def __aenter__(self) -> None:83        pass84 85    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:86        return False87 88 89async def tee_peer(90    iterator: AsyncIterator[T],91    # the buffer specific to this peer92    buffer: deque[T],93    # the buffers of all peers, including our own94    peers: list[deque[T]],95    lock: AbstractAsyncContextManager[Any],96) -> AsyncGenerator[T, None]:97    """Iterate over :py:func:`~.tee`."""98    try:99        while True:100            if not buffer:101                async with lock:102                    # Another peer produced an item while we were waiting for the lock.103                    # Proceed with the next loop iteration to yield the item.104                    if buffer:105                        continue106                    try:107                        item = await iterator.__anext__()108                    except StopAsyncIteration:109                        break110                    else:111                        # Append to all buffers, including our own. We'll fetch our112                        # item from the buffer again, instead of yielding it directly.113                        # This ensures the proper item ordering if any of our peers114                        # are fetching items concurrently. They may have buffered their115                        # item already.116                        for peer_buffer in peers:117                            peer_buffer.append(item)118            yield buffer.popleft()119    finally:120        async with lock:121            # this peer is done – remove its buffer122            for idx, peer_buffer in enumerate(peers):  # pragma: no branch123                if peer_buffer is buffer:124                    peers.pop(idx)125                    break126            # if we are the last peer, try and close the iterator127            if not peers and hasattr(iterator, "aclose"):128                await iterator.aclose()129 130 131class Tee(Generic[T]):132    """Create ``n`` separate asynchronous iterators over ``iterable``.133 134    This splits a single ``iterable`` into multiple iterators, each providing135    the same items in the same order.136    All child iterators may advance separately but pare the same items137    from ``iterable`` -- when the most advanced iterator retrieves an item,138    it is buffered until the least advanced iterator has yielded it as well.139    A ``tee`` works lazily and can handle an infinite ``iterable``, provided140    that all iterators advance.141 142    ```python143    async def derivative(sensor_data):144        previous, current = a.tee(sensor_data, n=2)145        await a.anext(previous)  # advance one iterator146        return a.map(operator.sub, previous, current)147    ```148 149    Unlike :py:func:`itertools.tee`, :py:func:`~.tee` returns a custom type instead150    of a :py:class:`tuple`. Like a tuple, it can be indexed, iterated and unpacked151    to get the child iterators. In addition, its :py:meth:`~.tee.aclose` method152    immediately closes all children, and it can be used in an ``async with`` context153    for the same effect.154 155    If ``iterable`` is an iterator and read elsewhere, ``tee`` will *not*156    provide these items. Also, ``tee`` must internally buffer each item until the157    last iterator has yielded it; if the most and least advanced iterator differ158    by most data, using a :py:class:`list` is more efficient (but not lazy).159 160    If the underlying iterable is concurrency safe (``anext`` may be awaited161    concurrently) the resulting iterators are concurrency safe as well. Otherwise,162    the iterators are safe if there is only ever one single "most advanced" iterator.163    To enforce sequential use of ``anext``, provide a ``lock``164    - e.g. an :py:class:`asyncio.Lock` instance in an :py:mod:`asyncio` application -165    and access is automatically synchronised.166    """167 168    def __init__(169        self,170        iterable: AsyncIterator[T],171        n: int = 2,172        *,173        lock: Optional[AbstractAsyncContextManager[Any]] = None,174    ):175        self._iterator = iterable.__aiter__()  # before 3.10 aiter() doesn't exist176        self._buffers: list[deque[T]] = [deque() for _ in range(n)]177        self._children = tuple(178            tee_peer(179                iterator=self._iterator,180                buffer=buffer,181                peers=self._buffers,182                lock=lock if lock is not None else NoLock(),183            )184            for buffer in self._buffers185        )186 187    def __len__(self) -> int:188        return len(self._children)189 190    @overload191    def __getitem__(self, item: int) -> AsyncIterator[T]: ...192 193    @overload194    def __getitem__(self, item: slice) -> tuple[AsyncIterator[T], ...]: ...195 196    def __getitem__(197        self, item: Union[int, slice]198    ) -> Union[AsyncIterator[T], tuple[AsyncIterator[T], ...]]:199        return self._children[item]200 201    def __iter__(self) -> Iterator[AsyncIterator[T]]:202        yield from self._children203 204    async def __aenter__(self) -> Tee[T]:205        return self206 207    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:208        await self.aclose()209        return False210 211    async def aclose(self) -> None:212        for child in self._children:213            await child.aclose()214 215 216atee = Tee217 218 219async def async_zip(*async_iterables):220    """Async version of zip."""221    # Before Python 3.10, aiter() was not available222    iterators = [iterable.__aiter__() for iterable in async_iterables]223    while True:224        try:225            items = await asyncio.gather(226                *(py_anext(iterator) for iterator in iterators)227            )228            yield tuple(items)229        except StopAsyncIteration:230            break231 232 233def ensure_async_iterator(234    iterable: Union[Iterable, AsyncIterable],235) -> AsyncIterator:236    if hasattr(iterable, "__anext__"):237        return cast(AsyncIterator, iterable)238    elif hasattr(iterable, "__aiter__"):239        return cast(AsyncIterator, iterable.__aiter__())240    else:241 242        class AsyncIteratorWrapper:243            def __init__(self, iterable: Iterable):244                self._iterator = iter(iterable)245 246            async def __anext__(self):247                try:248                    return next(self._iterator)249                except StopIteration:250                    raise StopAsyncIteration251 252            def __aiter__(self):253                return self254 255        return AsyncIteratorWrapper(iterable)256 257 258def aiter_with_concurrency(259    n: Optional[int],260    generator: AsyncIterator[Coroutine[None, None, T]],261    *,262    _eager_consumption_timeout: float = 0,263) -> AsyncGenerator[T, None]:264    """Process async generator with max parallelism.265 266    Args:267        n: The number of tasks to run concurrently.268        generator: The async generator to process.269        _eager_consumption_timeout: If set, check for completed tasks after270            each iteration and yield their results. This can be used to271            consume the generator eagerly while still respecting the concurrency272            limit.273 274    Yields:275        The processed items yielded by the async generator.276    """277    if n == 0:278 279        async def consume():280            async for item in generator:281                yield await item282 283        return consume()284    semaphore = cast(285        asyncio.Semaphore, asyncio.Semaphore(n) if n is not None else NoLock()286    )287 288    async def process_item(ix: int, item):289        async with semaphore:290            res = await item291            return (ix, res)292 293    async def process_generator():294        tasks = {}295        accepts_context = asyncio_accepts_context()296        ix = 0297        async for item in generator:298            if accepts_context:299                context = contextvars.copy_context()300                task = asyncio.create_task(process_item(ix, item), context=context)301            else:302                task = asyncio.create_task(process_item(ix, item))303            tasks[ix] = task304            ix += 1305            if _eager_consumption_timeout > 0:306                try:307                    for _fut in asyncio.as_completed(308                        tasks.values(),309                        timeout=_eager_consumption_timeout,310                    ):311                        task_idx, res = await _fut312                        yield res313                        del tasks[task_idx]314                except asyncio.TimeoutError:315                    pass316            if n is not None and len(tasks) >= n:317                done, _ = await asyncio.wait(318                    tasks.values(), return_when=asyncio.FIRST_COMPLETED319                )320                for task in done:321                    task_idx, res = task.result()322                    yield res323                    del tasks[task_idx]324 325        for task in asyncio.as_completed(tasks.values()):326            _, res = await task327            yield res328 329    return process_generator()330 331 332def accepts_context(callable: Callable[..., Any]) -> bool:333    """Check if a callable accepts a context argument."""334    try:335        return inspect.signature(callable).parameters.get("context") is not None336    except ValueError:337        return False338 339 340# Ported from Python 3.9+ to support Python 3.8341async def aio_to_thread(342    ctx: contextvars.Context,343    func,344    /,345    *args,346    **kwargs,347):348    """Run ``func`` in a separate thread, inside ``ctx``.349 350    ``ctx`` is the :class:`~contextvars.Context` in which ``func`` is invoked.351    Callers that want default isolation should pass352    ``contextvars.copy_context()``; callers with a specific Context353    (e.g. :func:`trace`) pass it directly so subsequent reads from that354    Context see the mutations.355 356    Return a coroutine that can be awaited to get the eventual result of ``func``.357    """358    overrides = get_runtime_overrides()359    if overrides.aio_to_thread is not None:360        return await overrides.aio_to_thread(361            _default_aio_to_thread, ctx, func, *args, **kwargs362        )363    return await _default_aio_to_thread(ctx, func, *args, **kwargs)364 365 366async def _default_aio_to_thread(367    ctx: contextvars.Context,368    func,369    /,370    *args,371    **kwargs,372):373    """Default implementation of aio_to_thread using run_in_executor."""374    loop = asyncio.get_running_loop()375    func_call = functools.partial(ctx.run, func, *args, **kwargs)376    return await loop.run_in_executor(None, func_call)377 378 379@functools.lru_cache(maxsize=1)380def asyncio_accepts_context():381    """Check if the current asyncio event loop accepts a context argument."""382    return accepts_context(asyncio.create_task)383