codekingpro/portable-devtools

1"""Utilities for working with iterators."""2 3from collections import deque4from collections.abc import Generator, Iterable, Iterator5from contextlib import AbstractContextManager6from itertools import islice7from types import TracebackType8from typing import (9    Any,10    Generic,11    Literal,12    TypeVar,13    overload,14)15 16T = TypeVar("T")17 18 19class NoLock:20    """Dummy lock that provides the proper interface but no protection."""21 22    def __enter__(self) -> None:23        """Do nothing."""24 25    def __exit__(26        self,27        exc_type: type[BaseException] | None,28        exc_val: BaseException | None,29        exc_tb: TracebackType | None,30    ) -> Literal[False]:31        """Return False (exception not suppressed)."""32        return False33 34 35def tee_peer(36    iterator: Iterator[T],37    # the buffer specific to this peer38    buffer: deque[T],39    # the buffers of all peers, including our own40    peers: list[deque[T]],41    lock: AbstractContextManager[Any],42) -> Generator[T, None, None]:43    """An individual iterator of a `.tee`.44 45    This function is a generator that yields items from the shared iterator `iterator`.46    It buffers items until the least advanced iterator has yielded them as well. The47    buffer is shared with all other peers.48 49    Args:50        iterator: The shared iterator.51        buffer: The buffer for this peer.52        peers: The buffers of all peers.53        lock: The lock to synchronise access to the shared buffers.54 55    Yields:56        The next item from the shared iterator.57    """58    try:59        while True:60            if not buffer:61                with lock:62                    # Another peer produced an item while we were waiting for the lock.63                    # Proceed with the next loop iteration to yield the item.64                    if buffer:65                        continue66                    try:67                        item = next(iterator)68                    except StopIteration:69                        break70                    else:71                        # Append to all buffers, including our own. We'll fetch our72                        # item from the buffer again, instead of yielding it directly.73                        # This ensures the proper item ordering if any of our peers74                        # are fetching items concurrently. They may have buffered their75                        # item already.76                        for peer_buffer in peers:77                            peer_buffer.append(item)78            yield buffer.popleft()79    finally:80        with lock:81            # this peer is done - remove its buffer82            for idx, peer_buffer in enumerate(peers):  # pragma: no branch83                if peer_buffer is buffer:84                    peers.pop(idx)85                    break86            # if we are the last peer, try and close the iterator87            if not peers and hasattr(iterator, "close"):88                iterator.close()89 90 91class Tee(Generic[T]):92    """Create `n` separate asynchronous iterators over `iterable`.93 94    This splits a single `iterable` into multiple iterators, each providing the same95    items in the same order.96 97    All child iterators may advance separately but share the same items from `iterable`98    -- when the most advanced iterator retrieves an item, it is buffered until the least99    advanced iterator has yielded it as well. A `tee` works lazily and can handle an100    infinite `iterable`, provided that all iterators advance.101 102    ```python103    async def derivative(sensor_data):104        previous, current = a.tee(sensor_data, n=2)105        await a.anext(previous)  # advance one iterator106        return a.map(operator.sub, previous, current)107    ```108 109    Unlike `itertools.tee`, `.tee` returns a custom type instead of a `tuple`. Like a110    tuple, it can be indexed, iterated and unpacked to get the child iterators. In111    addition, its `.tee.aclose` method immediately closes all children, and it can be112    used in an `async with` context for the same effect.113 114    If `iterable` is an iterator and read elsewhere, `tee` will *not* provide these115    items. Also, `tee` must internally buffer each item until the last iterator has116    yielded it; if the most and least advanced iterator differ by most data, using a117    `list` is more efficient (but not lazy).118 119    If the underlying iterable is concurrency safe (`anext` may be awaited concurrently)120    the resulting iterators are concurrency safe as well. Otherwise, the iterators are121    safe if there is only ever one single "most advanced" iterator. To enforce122    sequential use of `anext`, provide a `lock`123 124    - e.g., an `asyncio.Lock` instance in an `asyncio` application - and access is125        automatically synchronised.126 127    """128 129    def __init__(130        self,131        iterable: Iterator[T],132        n: int = 2,133        *,134        lock: AbstractContextManager[Any] | None = None,135    ):136        """Create a `tee`.137 138        Args:139            iterable: The iterable to split.140            n: The number of iterators to create.141            lock: The lock to synchronise access to the shared buffers.142 143        """144        self._iterator = iter(iterable)145        self._buffers: list[deque[T]] = [deque() for _ in range(n)]146        self._children = tuple(147            tee_peer(148                iterator=self._iterator,149                buffer=buffer,150                peers=self._buffers,151                lock=lock if lock is not None else NoLock(),152            )153            for buffer in self._buffers154        )155 156    def __len__(self) -> int:157        """Return the number of child iterators."""158        return len(self._children)159 160    @overload161    def __getitem__(self, item: int) -> Iterator[T]: ...162 163    @overload164    def __getitem__(self, item: slice) -> tuple[Iterator[T], ...]: ...165 166    def __getitem__(self, item: int | slice) -> Iterator[T] | tuple[Iterator[T], ...]:167        """Return the child iterator(s) at the given index or slice."""168        return self._children[item]169 170    def __iter__(self) -> Iterator[Iterator[T]]:171        """Return an iterator over the child iterators.172 173        Yields:174            The child iterators.175        """176        yield from self._children177 178    def __enter__(self) -> "Tee[T]":179        """Return `Tee` instance."""180        return self181 182    def __exit__(183        self,184        exc_type: type[BaseException] | None,185        exc_val: BaseException | None,186        exc_tb: TracebackType | None,187    ) -> Literal[False]:188        """Close all child iterators.189 190        Returns:191            `False` (exception not suppressed).192        """193        self.close()194        return False195 196    def close(self) -> None:197        """Close all child iterators."""198        for child in self._children:199            child.close()200 201 202# Why this is needed https://stackoverflow.com/a/44638570203safetee = Tee204 205 206def batch_iterate(size: int | None, iterable: Iterable[T]) -> Iterator[list[T]]:207    """Utility batching function.208 209    Args:210        size: The size of the batch.211 212            If `None`, returns a single batch.213        iterable: The iterable to batch.214 215    Yields:216        The batches of the iterable.217    """218    it = iter(iterable)219    while True:220        chunk = list(islice(it, size))221        if not chunk:222            return223        yield chunk224