iter.py — langchain Source File

Source Code

"""Utilities for working with iterators."""

from collections import deque
from collections.abc import Generator, Iterable, Iterator
from contextlib import AbstractContextManager
from itertools import islice
from types import TracebackType
from typing import (
    Any,
    Generic,
    Literal,
    TypeVar,
    overload,
)

T = TypeVar("T")


class NoLock:
    """Dummy lock that provides the proper interface but no protection."""

    def __enter__(self) -> None:
        """Do nothing."""

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> Literal[False]:
        """Return False (exception not suppressed)."""
        return False


def tee_peer(
    iterator: Iterator[T],
    # the buffer specific to this peer
    buffer: deque[T],
    # the buffers of all peers, including our own
    peers: list[deque[T]],
    lock: AbstractContextManager[Any],
) -> Generator[T, None, None]:
    """An individual iterator of a `.tee`.

    This function is a generator that yields items from the shared iterator `iterator`.
    It buffers items until the least advanced iterator has yielded them as well. The
    buffer is shared with all other peers.

    Args:
        iterator: The shared iterator.
        buffer: The buffer for this peer.
        peers: The buffers of all peers.
        lock: The lock to synchronise access to the shared buffers.

    Yields:
        The next item from the shared iterator.
    """
    try:
        while True:
            if not buffer:
// ... (164 more lines)