admin / Synapse-Cortex

from __future__ import annotations

__all__ = (
    "accumulate",
    "batched",
    "Chain",
    "combinations",
    "combinations_with_replacement",
    "compress",
    "count",
    "cycle",
    "dropwhile",
    "filterfalse",
    "groupby",
    "islice",
    "pairwise",
    "permutations",
    "product",
    "repeat",
    "starmap",
    "tee",
    "takewhile",
    "zip_longest",
)

import itertools
import operator
import sys
from collections.abc import (
    AsyncGenerator,
    AsyncIterable,
    AsyncIterator,
    Awaitable,
    Callable,
    Iterable,
    Iterator,
)
from dataclasses import dataclass, field
from typing import Any, Generic, TypeVar, cast, overload

from ._core._synchronization import Lock
from ._core._tasks import CancelScope
from .lowlevel import cancel_shielded_checkpoint, checkpoint, checkpoint_if_cancelled

T = TypeVar("T")
R = TypeVar("R")
_tee_end = object()


@dataclass(eq=False)
class _IterableAsyncIterator(AsyncIterator[T]):
    iterator: Iterator[T]

    async def __anext__(self) -> T:
        await checkpoint_if_cancelled()
        try:
            result = next(self.iterator)
        except StopIteration:
            await cancel_shielded_checkpoint()
            raise StopAsyncIteration from None

        await cancel_shielded_checkpoint()
        return result


def _iterate(iterable: Iterable[T] | AsyncIterable[T]) -> AsyncIterator[T]:
    if isinstance(iterable, AsyncIterator):
        return iterable

    if isinstance(iterable, AsyncIterable):
        return iterable.__aiter__()

    return _IterableAsyncIterator(iter(iterable))


@dataclass(eq=False)
class _TeeLink(Generic[T]):
    value: object | None = None
    next: _TeeLink[T] | None = None
    filled: bool = False


@dataclass(eq=False)
class _TeeState(Generic[T]):
    iterator: AsyncIterator[T]
    lock: Lock = field(default_factory=Lock)

    async def fill(self, link: _TeeLink[T]) -> bool:
        if link.filled:
            return False

        async with self.lock:
            if link.filled:
                return True

            link.value = await anext(self.iterator, _tee_end)
            if link.value is not _tee_end:
                link.next = _TeeLink()

            link.filled = True
            return True


class _TeeAsyncIterator(AsyncIterator[T]):
    _state: _TeeState[T]
    _link: _TeeLink[T]
    _element_yielded: bool

    def __init__(
        self, iterable: Iterable[T] | AsyncIterable[T] | _TeeAsyncIterator[T]
    ) -> None:
        if isinstance(iterable, _TeeAsyncIterator):
            self._state = iterable._state
            self._link = iterable._link
        else:
            self._state = _TeeState(_iterate(iterable))
            self._link = _TeeLink()

        self._element_yielded = False

    async def __anext__(self) -> T:
        had_yieldpoint = await self._state.fill(self._link)
        if self._link.value is _tee_end:
            if not self._element_yielded:
                await checkpoint()

            raise StopAsyncIteration

        if not had_yieldpoint:
            await checkpoint_if_cancelled()

        self._element_yielded = True
        value = cast(T, self._link.value)
        next_link = self._link.next
        assert next_link is not None
        self._link = next_link
        if not had_yieldpoint:
            await cancel_shielded_checkpoint()

        return value


async def _operator_add(x: T, y: T) -> T:
    return operator.add(x, y)


async def accumulate(
    iterable: Iterable[T] | AsyncIterable[T],
    function: Callable[[T, T], Awaitable[T]] = _operator_add,
    *,
    initial: T | None = None,
) -> AsyncGenerator[T, None]:
    iterator = _iterate(iterable)
    if initial is None:
        try:
            total = await anext(iterator)
        except StopAsyncIteration:
            await checkpoint()
            return
    else:
        await checkpoint_if_cancelled()
        total = initial
        await cancel_shielded_checkpoint()

    yield total

    async for element in iterator:
        total = await function(total, element)
        yield total


async def batched(
    iterable: Iterable[T] | AsyncIterable[T], n: int, *, strict: bool = False
) -> AsyncGenerator[tuple[T, ...], None]:
    if n < 1:
        raise ValueError("n must be at least one")

    iterator = _iterate(iterable)

    while True:
        batch: list[T] = []
        for _ in range(n):
            try:
                batch.append(await anext(iterator))
            except StopAsyncIteration:
                if not batch:
                    await checkpoint()
                    return
                if strict:
                    raise ValueError("batched(): incomplete batch") from None

                yield tuple(batch)
                return

        yield tuple(batch)


class Chain:
    def __call__(
        self, *iterables: Iterable[T] | AsyncIterable[T]
    ) -> AsyncGenerator[T, None]:
        return self.from_iterable(iterables)

    async def from_iterable(
        self,
        iterables: (
            Iterable[Iterable[T] | AsyncIterable[T]]
            | AsyncIterable[Iterable[T] | AsyncIterable[T]]
        ),
    ) -> AsyncGenerator[T, None]:
        element_yielded = False
        outer_iter = _iterate(iterables)

        try:
            async for iterable in outer_iter:
                async for element in _iterate(iterable):
                    element_yielded = True
                    yield element
        finally:
            aclose = getattr(outer_iter, "aclose", None)
            if aclose is not None:
                with CancelScope(shield=True):
                    await aclose()

        if not element_yielded:
            await checkpoint()


chain: Chain = Chain()


async def combinations(
    iterable: Iterable[T] | AsyncIterable[T], r: int
) -> AsyncGenerator[tuple[T, ...], None]:
    pool: list[T] = [element async for element in _iterate(iterable)]
    async for combination in _iterate(itertools.combinations(pool, r)):
        yield combination


async def combinations_with_replacement(
    iterable: Iterable[T] | AsyncIterable[T], r: int
) -> AsyncGenerator[tuple[T, ...], None]:
    pool: list[T] = [element async for element in _iterate(iterable)]
    async for combination in _iterate(itertools.combinations_with_replacement(pool, r)):
        yield combination


async def compress(
    data: Iterable[T] | AsyncIterable[T],
    selectors: Iterable[object] | AsyncIterable[object],
) -> AsyncGenerator[T, None]:
    data_iterator = _iterate(data)
    selector_iterator = _iterate(selectors)
    element_yielded = False

    while True:
        try:
            datum = await anext(data_iterator)
            selector = await anext(selector_iterator)
        except StopAsyncIteration:
            if not element_yielded:
                await checkpoint()

            return

        if selector:
            element_yielded = True
            yield datum


async def count(start: int = 0, step: int = 1) -> AsyncGenerator[int, None]:
    n = start
    while True:
        await checkpoint_if_cancelled()
        value = n
        n += step
        await cancel_shielded_checkpoint()
        yield value


async def cycle(
    iterable: Iterable[T] | AsyncIterable[T],
) -> AsyncGenerator[T, None]:
    saved: list[T] = []
    async for element in _iterate(iterable):
        saved.append(element)
        yield element

    if not saved:
        await checkpoint()
        return

    while True:
        for element in saved:
            await checkpoint()
            yield element


async def dropwhile(
    predicate: Callable[[T], Awaitable[object]],
    iterable: Iterable[T] | AsyncIterable[T],
) -> AsyncGenerator[T, None]:
    element_yielded = False
    dropping = True

    async for element in _iterate(iterable):
        if dropping and await predicate(element):
            continue

        dropping = False
        element_yielded = True
        yield element

    if not element_yielded:
        await checkpoint()


async def filterfalse(
    predicate: Callable[[T], Awaitable[object]],
    iterable: Iterable[T] | AsyncIterable[T],
) -> AsyncGenerator[T, None]:
    element_yielded = False

    async for element in _iterate(iterable):
        if not await predicate(element):
            element_yielded = True
            yield element

    if not element_yielded:
        await checkpoint()


@overload
def groupby(
    iterable: Iterable[T] | AsyncIterable[T],
) -> AsyncGenerator[tuple[T, list[T]], None]: ...


@overload
def groupby(
    iterable: Iterable[T] | AsyncIterable[T],
    key: Callable[[T], Awaitable[R]],
) -> AsyncGenerator[tuple[R, list[T]], None]: ...


async def groupby(
    iterable: Iterable[T] | AsyncIterable[T],
    key: Callable[[T], Awaitable[object]] | None = None,
) -> AsyncGenerator[tuple[object, list[T]], None]:
    iterator = _iterate(iterable)
    try:
        element = await anext(iterator)
    except StopAsyncIteration:
        await checkpoint()
        return

    group_key = element if key is None else await key(element)
    values = [element]

    async for element in iterator:
        next_key = element if key is None else await key(element)
        if next_key != group_key:
            completed_group = group_key, values
            group_key = next_key
            values = [element]
            yield completed_group
        else:
            values.append(element)

    yield group_key, values


@overload
def islice(
    iterable: Iterable[T] | AsyncIterable[T],
    stop: int | None,
    /,
) -> AsyncGenerator[T, None]: ...


@overload
def islice(
    iterable: Iterable[T] | AsyncIterable[T],
    start: int | None,
    stop: int | None,
    step: int | None = 1,
    /,
) -> AsyncGenerator[T, None]: ...


async def islice(
    iterable: Iterable[T] | AsyncIterable[T],
    *args: int | None,
) -> AsyncGenerator[T, None]:
    if not args:
        raise TypeError("islice expected at least 2 arguments, got 1")
    if len(args) > 3:
        raise TypeError(f"islice expected at most 4 arguments, got {len(args) + 1}")

    slice_args = slice(*args)

    start_message = (
        "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize."
    )
    stop_message = (
        "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize."
    )
    step_message = "Step for islice() must be a positive integer or None."

    def normalize_index(value: object, message: str) -> int:
        try:
            index = operator.index(cast(Any, value))
        except TypeError:
            raise ValueError(message) from None

        if index < 0 or index > sys.maxsize:
            raise ValueError(message)

        return index

    start = (
        0
        if slice_args.start is None
        else normalize_index(slice_args.start, start_message)
    )
    stop = (
        None
        if slice_args.stop is None
        else normalize_index(slice_args.stop, stop_message)
    )
    step = (
        1 if slice_args.step is None else normalize_index(slice_args.step, step_message)
    )

    if step <= 0:
        raise ValueError(step_message)

    if stop == 0 or start == stop:
        await checkpoint()
        return

    iterator = _iterate(iterable)
    index = 0
    element_yielded = False

    while stop is None or index < stop:
        try:
            element = await anext(iterator)
        except StopAsyncIteration:
            if not element_yielded:
                await checkpoint()

            return

        if index >= start and (index - start) % step == 0:
            index += 1
            element_yielded = True
            yield element
        else:
            index += 1

    if not element_yielded:
        await checkpoint()


async def pairwise(
    iterable: Iterable[T] | AsyncIterable[T],
) -> AsyncGenerator[tuple[T, T], None]:
    iterator = _iterate(iterable)
    try:
        previous = await anext(iterator)
    except StopAsyncIteration:
        await checkpoint()
        return

    element_yielded = False
    async for element in iterator:
        element_yielded = True
        pair = (previous, element)
        previous = element
        yield pair

    if not element_yielded:
        await checkpoint()


async def permutations(
    iterable: Iterable[T] | AsyncIterable[T], r: int | None = None
) -> AsyncGenerator[tuple[T, ...], None]:
    pool: list[T] = [element async for element in _iterate(iterable)]
    n = len(pool)
    if r is None:
        r = n
    elif not isinstance(r, int):
        raise TypeError("Expected int as r")
    elif r < 0:
        raise ValueError("r must be non-negative")

    async for permutation in _iterate(itertools.permutations(pool, r)):
        yield permutation


async def product(
    *iterables: Iterable[T] | AsyncIterable[T], repeat: int = 1
) -> AsyncGenerator[tuple[T, ...], None]:
    repeat = operator.index(repeat)
    if repeat < 0:
        raise ValueError("repeat argument cannot be negative")

    pools: list[tuple[T, ...]] = []
    for iterable in iterables:
        pool: list[T] = [element async for element in _iterate(iterable)]
        pools.append(tuple(pool))

    async for value in _iterate(itertools.product(*pools, repeat=repeat)):
        yield value


async def repeat(element: T, times: int | None = None) -> AsyncGenerator[T, None]:
    if times is None:
        while True:
            await checkpoint()
            yield element

    remaining = operator.index(cast(Any, times))
    if remaining <= 0:
        await checkpoint()
        return

    while remaining > 0:
        await checkpoint_if_cancelled()
        remaining -= 1
        await cancel_shielded_checkpoint()
        yield element


async def starmap(
    function: Callable[..., Awaitable[R]],
    iterable: (
        Iterable[Iterable[object] | AsyncIterable[object]]
        | AsyncIterable[Iterable[object] | AsyncIterable[object]]
    ),
) -> AsyncGenerator[R, None]:
    result_yielded = False

    async for args_iterable in _iterate(iterable):
        args = [element async for element in _iterate(args_iterable)]
        result_yielded = True
        yield await function(*args)

    if not result_yielded:
        await checkpoint()


def tee(
    iterable: Iterable[T] | AsyncIterable[T], n: int = 2
) -> tuple[AsyncIterator[T], ...]:
    n = operator.index(cast(Any, n))
    if n < 0:
        raise ValueError("n must be >= 0")
    if n == 0:
        return ()

    iterator = _TeeAsyncIterator(iterable)
    iterators: list[AsyncIterator[T]] = [iterator]
    iterators.extend(_TeeAsyncIterator(iterator) for _ in range(n - 1))
    return tuple(iterators)


async def takewhile(
    predicate: Callable[[T], Awaitable[object]],
    iterable: Iterable[T] | AsyncIterable[T],
) -> AsyncGenerator[T, None]:
    element_yielded = False

    async for element in _iterate(iterable):
        if not await predicate(element):
            if not element_yielded:
                await checkpoint()

            return

        element_yielded = True
        yield element

    if not element_yielded:
        await checkpoint()


async def zip_longest(
    *iterables: Iterable[object] | AsyncIterable[object],
    fillvalue: object = None,
) -> AsyncGenerator[tuple[object, ...], None]:
    iterators = [_iterate(iterable) for iterable in iterables]
    num_active = len(iterators)
    if not num_active:
        await checkpoint()
        return

    active = [True] * num_active
    tuple_yielded = False

    while True:
        values: list[object] = []
        for index, iterator in enumerate(iterators):
            if not active[index]:
                values.append(fillvalue)
                continue

            try:
                value = await anext(iterator)
            except StopAsyncIteration:
                active[index] = False
                num_active -= 1
                if not num_active:
                    if not tuple_yielded:
                        await checkpoint()

                    return

                value = fillvalue

            values.append(value)

        tuple_yielded = True
        yield tuple(values)