1 files changed, 284 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/aiostream/stream/select.py b/.venv/lib/python3.12/site-packages/aiostream/stream/select.py
new file mode 100644
index 00000000..9390f464
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/aiostream/stream/select.py
@@ -0,0 +1,284 @@
+"""Selection operators."""
+from __future__ import annotations
+
+import asyncio
+import builtins
+import collections
+
+from typing import Awaitable, Callable, TypeVar, AsyncIterable, AsyncIterator
+
+from . import transform
+from ..aiter_utils import aiter, anext
+from ..core import streamcontext, pipable_operator
+
+__all__ = [
+    "take",
+    "takelast",
+    "skip",
+    "skiplast",
+    "getitem",
+    "filter",
+    "until",
+    "dropwhile",
+    "takewhile",
+]
+
+T = TypeVar("T")
+
+
+@pipable_operator
+async def take(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
+    """Forward the first ``n`` elements from an asynchronous sequence.
+
+    If ``n`` is negative, it simply terminates before iterating the source.
+    """
+    enumerated = transform.enumerate.raw(source)
+    async with streamcontext(enumerated) as streamer:
+        if n <= 0:
+            return
+        async for i, item in streamer:
+            yield item
+            if i >= n - 1:
+                return
+
+
+@pipable_operator
+async def takelast(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
+    """Forward the last ``n`` elements from an asynchronous sequence.
+
+    If ``n`` is negative, it simply terminates after iterating the source.
+
+    Note: it is required to reach the end of the source before the first
+    element is generated.
+    """
+    queue: collections.deque[T] = collections.deque(maxlen=n if n > 0 else 0)
+    async with streamcontext(source) as streamer:
+        async for item in streamer:
+            queue.append(item)
+        for item in queue:
+            yield item
+
+
+@pipable_operator
+async def skip(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
+    """Forward an asynchronous sequence, skipping the first ``n`` elements.
+
+    If ``n`` is negative, no elements are skipped.
+    """
+    enumerated = transform.enumerate.raw(source)
+    async with streamcontext(enumerated) as streamer:
+        async for i, item in streamer:
+            if i >= n:
+                yield item
+
+
+@pipable_operator
+async def skiplast(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
+    """Forward an asynchronous sequence, skipping the last ``n`` elements.
+
+    If ``n`` is negative, no elements are skipped.
+
+    Note: it is required to reach the ``n+1`` th element of the source
+    before the first element is generated.
+    """
+    queue: collections.deque[T] = collections.deque(maxlen=n if n > 0 else 0)
+    async with streamcontext(source) as streamer:
+        async for item in streamer:
+            if n <= 0:
+                yield item
+                continue
+            if len(queue) == n:
+                yield queue[0]
+            queue.append(item)
+
+
+@pipable_operator
+async def filterindex(
+    source: AsyncIterable[T], func: Callable[[int], bool]
+) -> AsyncIterator[T]:
+    """Filter an asynchronous sequence using the index of the elements.
+
+    The given function is synchronous, takes the index as an argument,
+    and returns ``True`` if the corresponding should be forwarded,
+    ``False`` otherwise.
+    """
+    enumerated = transform.enumerate.raw(source)
+    async with streamcontext(enumerated) as streamer:
+        async for i, item in streamer:
+            if func(i):
+                yield item
+
+
+@pipable_operator
+def slice(source: AsyncIterable[T], *args: int) -> AsyncIterator[T]:
+    """Slice an asynchronous sequence.
+
+    The arguments are the same as the builtin type slice.
+
+    There are two limitations compare to regular slices:
+    - Positive stop index with negative start index is not supported
+    - Negative step is not supported
+    """
+    s = builtins.slice(*args)
+    start, stop, step = s.start or 0, s.stop, s.step or 1
+    aiterator = aiter(source)
+    # Filter the first items
+    if start < 0:
+        aiterator = takelast.raw(aiterator, abs(start))
+    elif start > 0:
+        aiterator = skip.raw(aiterator, start)
+    # Filter the last items
+    if stop is not None:
+        if stop >= 0 and start < 0:
+            raise ValueError("Positive stop with negative start is not supported")
+        elif stop >= 0:
+            aiterator = take.raw(aiterator, stop - start)
+        else:
+            aiterator = skiplast.raw(aiterator, abs(stop))
+    # Filter step items
+    if step is not None:
+        if step > 1:
+            aiterator = filterindex.raw(aiterator, lambda i: i % step == 0)
+        elif step < 0:
+            raise ValueError("Negative step not supported")
+    # Return
+    return aiterator
+
+
+@pipable_operator
+async def item(source: AsyncIterable[T], index: int) -> AsyncIterator[T]:
+    """Forward the ``n``th element of an asynchronous sequence.
+
+    The index can be negative and works like regular indexing.
+    If the index is out of range, and ``IndexError`` is raised.
+    """
+    # Prepare
+    if index >= 0:
+        source = skip.raw(source, index)
+    else:
+        source = takelast(source, abs(index))
+    async with streamcontext(source) as streamer:
+        # Get first item
+        try:
+            result = await anext(streamer)
+        except StopAsyncIteration:
+            raise IndexError("Index out of range")
+        # Check length
+        if index < 0:
+            count = 1
+            async for _ in streamer:
+                count += 1
+            if count != abs(index):
+                raise IndexError("Index out of range")
+        # Yield result
+        yield result
+
+
+@pipable_operator
+def getitem(source: AsyncIterable[T], index: int | builtins.slice) -> AsyncIterator[T]:
+    """Forward one or several items from an asynchronous sequence.
+
+    The argument can either be a slice or an integer.
+    See the slice and item operators for more information.
+    """
+    if isinstance(index, builtins.slice):
+        return slice.raw(source, index.start, index.stop, index.step)
+    if isinstance(index, int):
+        return item.raw(source, index)
+    raise TypeError("Not a valid index (int or slice)")
+
+
+@pipable_operator
+async def filter(
+    source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
+) -> AsyncIterator[T]:
+    """Filter an asynchronous sequence using an arbitrary function.
+
+    The function takes the item as an argument and returns ``True``
+    if it should be forwarded, ``False`` otherwise.
+    The function can either be synchronous or asynchronous.
+    """
+    iscorofunc = asyncio.iscoroutinefunction(func)
+    async with streamcontext(source) as streamer:
+        async for item in streamer:
+            result = func(item)
+            if iscorofunc:
+                assert isinstance(result, Awaitable)
+                result = await result
+            if result:
+                yield item
+
+
+@pipable_operator
+async def until(
+    source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
+) -> AsyncIterator[T]:
+    """Forward an asynchronous sequence until a condition is met.
+
+    Contrary to the ``takewhile`` operator, the last tested element is included
+    in the sequence.
+
+    The given function takes the item as an argument and returns a boolean
+    corresponding to the condition to meet. The function can either be
+    synchronous or asynchronous.
+    """
+    iscorofunc = asyncio.iscoroutinefunction(func)
+    async with streamcontext(source) as streamer:
+        async for item in streamer:
+            result = func(item)
+            if iscorofunc:
+                assert isinstance(result, Awaitable)
+                result = await result
+            yield item
+            if result:
+                return
+
+
+@pipable_operator
+async def takewhile(
+    source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
+) -> AsyncIterator[T]:
+    """Forward an asynchronous sequence while a condition is met.
+
+    Contrary to the ``until`` operator, the last tested element is not included
+    in the sequence.
+
+    The given function takes the item as an argument and returns a boolean
+    corresponding to the condition to meet. The function can either be
+    synchronous or asynchronous.
+    """
+    iscorofunc = asyncio.iscoroutinefunction(func)
+    async with streamcontext(source) as streamer:
+        async for item in streamer:
+            result = func(item)
+            if iscorofunc:
+                assert isinstance(result, Awaitable)
+                result = await result
+            if not result:
+                return
+            yield item
+
+
+@pipable_operator
+async def dropwhile(
+    source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
+) -> AsyncIterator[T]:
+    """Discard the elements from an asynchronous sequence
+    while a condition is met.
+
+    The given function takes the item as an argument and returns a boolean
+    corresponding to the condition to meet. The function can either be
+    synchronous or asynchronous.
+    """
+    iscorofunc = asyncio.iscoroutinefunction(func)
+    async with streamcontext(source) as streamer:
+        async for item in streamer:
+            result = func(item)
+            if iscorofunc:
+                assert isinstance(result, Awaitable)
+                result = await result
+            if not result:
+                yield item
+                break
+        async for item in streamer:
+            yield item