streamish 0.1.0__py3-none-any.whl

streamish/__init__.py

@@ -0,0 +1,63 @@
+"""Streamish - Iterator and async iterator utilities."""
+
+from collections.abc import AsyncIterable, Iterable
+
+from streamish.ops import (
+    batch,
+    chain,
+    chain_async,
+    distinct,
+    distinct_by,
+    enumerate,
+    filter,
+    flat_map,
+    flatten,
+    interleave,
+    map,
+    map_async,
+    merge,
+    partition,
+    partition_async,
+    scan,
+    skip,
+    skip_while,
+    take,
+    take_while,
+    window,
+    zip,
+    zip_async,
+)
+from streamish.stream import Stream
+
+__all__ = [
+    "stream",
+    "Stream",
+    "map",
+    "map_async",
+    "filter",
+    "take",
+    "skip",
+    "take_while",
+    "skip_while",
+    "distinct",
+    "distinct_by",
+    "flatten",
+    "flat_map",
+    "enumerate",
+    "scan",
+    "batch",
+    "window",
+    "partition",
+    "partition_async",
+    "zip",
+    "zip_async",
+    "chain",
+    "chain_async",
+    "interleave",
+    "merge",
+]
+
+
+def stream[T](source: Iterable[T] | AsyncIterable[T]) -> Stream[T]:
+    """Create a Stream from an iterable or async iterable."""
+    return Stream(source)

streamish/_util.py

@@ -0,0 +1,38 @@
+"""Internal utility helpers."""
+
+import inspect
+from collections.abc import (
+    AsyncGenerator,
+    AsyncIterable,
+    AsyncIterator,
+    Callable,
+    Iterable,
+    Iterator,
+)
+from typing import Any, TypeGuard
+
+__all__ = ["is_async_iterable", "is_awaitable", "ensure_async_iterator"]
+
+
+def is_async_iterable[T](
+    obj: Iterable[T] | AsyncIterable[T],
+) -> TypeGuard[AsyncIterable[T]]:
+    """Check if object is an async iterable."""
+    return hasattr(obj, "__aiter__")
+
+
+def is_awaitable(fn: Callable[..., Any]) -> bool:
+    """Check if function is async (returns awaitable)."""
+    return inspect.iscoroutinefunction(fn)
+
+
+async def ensure_async_iterator[T](
+    it: Iterator[T] | AsyncIterator[T],
+) -> AsyncGenerator[T]:
+    """Convert sync iterator to async iterator if needed."""
+    if isinstance(it, AsyncIterator):
+        async for item in it:
+            yield item
+    else:
+        for item in it:
+            yield item

streamish/ops/__init__.py

@@ -0,0 +1,43 @@
+"""Stream operations."""
+
+from streamish.ops.combine import chain, chain_async, interleave, merge, zip_async
+from streamish.ops.combine import zip_ as zip
+from streamish.ops.filter import (
+    distinct,
+    distinct_by,
+    skip,
+    skip_while,
+    take,
+    take_while,
+)
+from streamish.ops.group import batch, partition, partition_async, window
+from streamish.ops.transform import enumerate_ as enumerate
+from streamish.ops.transform import filter_ as filter
+from streamish.ops.transform import flat_map, flatten, map_async, scan
+from streamish.ops.transform import map_ as map
+
+__all__ = [
+    "map",
+    "filter",
+    "take",
+    "skip",
+    "take_while",
+    "skip_while",
+    "distinct",
+    "distinct_by",
+    "enumerate",
+    "scan",
+    "flatten",
+    "flat_map",
+    "batch",
+    "window",
+    "partition",
+    "partition_async",
+    "zip",
+    "zip_async",
+    "chain",
+    "chain_async",
+    "interleave",
+    "merge",
+    "map_async",
+]

streamish/ops/combine.py

@@ -0,0 +1,101 @@
+"""Combine operations."""
+
+import asyncio
+from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator
+
+from streamish._util import ensure_async_iterator, is_async_iterable
+
+__all__ = ["zip_", "zip_async", "chain", "chain_async", "interleave", "merge"]
+
+
+def zip_[T](*iterables: Iterable[T]) -> Iterator[tuple[T, ...]]:
+    """Zip iterables together."""
+    iters = [iter(it) for it in iterables]
+    while True:
+        result: list[T] = []
+        for it in iters:
+            try:
+                result.append(next(it))
+            except StopIteration:
+                return
+        yield tuple(result)
+
+
+async def zip_async[T](
+    *iterables: AsyncIterable[T] | Iterable[T],
+) -> AsyncIterator[tuple[T, ...]]:
+    """Zip async iterables together."""
+    aiters: list[AsyncIterator[T]] = [
+        it.__aiter__() if is_async_iterable(it) else ensure_async_iterator(iter(it))  # type: ignore[union-attr, arg-type]
+        for it in iterables
+    ]
+    while True:
+        try:
+            results: list[T] = []
+            for ait in aiters:
+                results.append(await ait.__anext__())
+            yield tuple(results)
+        except StopAsyncIteration:
+            break
+
+
+def chain[T](*iterables: Iterable[T]) -> Iterator[T]:
+    """Chain iterables together."""
+    for it in iterables:
+        yield from it
+
+
+async def chain_async[T](
+    *iterables: AsyncIterable[T] | Iterable[T],
+) -> AsyncIterator[T]:
+    """Chain async iterables together."""
+    for it in iterables:
+        if is_async_iterable(it):
+            async for item in it:  # type: ignore[union-attr]
+                yield item
+        else:
+            for item in it:  # type: ignore[union-attr]
+                yield item
+
+
+def interleave[T](*iterables: Iterable[T]) -> Iterator[T]:
+    """Alternate elements from iterables (round-robin)."""
+    iters: list[Iterator[T]] = [iter(it) for it in iterables]
+    while iters:
+        next_iters: list[Iterator[T]] = []
+        for it in iters:
+            try:
+                yield next(it)
+                next_iters.append(it)
+            except StopIteration:
+                pass
+        iters = next_iters
+
+
+async def merge[T](*iterables: AsyncIterable[T]) -> AsyncIterator[T]:
+    """Merge async iterables - emit as items arrive."""
+    pending: set[asyncio.Task[tuple[int, T | None, bool]]] = set()
+    aiters: dict[int, AsyncIterator[T]] = {}
+
+    for i, it in enumerate(iterables):
+        aiters[i] = it.__aiter__()
+        task = asyncio.create_task(_fetch_next(i, aiters[i]))
+        pending.add(task)
+
+    while pending:
+        done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+        for task in done:
+            idx, value, exhausted = task.result()
+            if exhausted:
+                continue
+            yield value  # type: ignore[misc]
+            new_task = asyncio.create_task(_fetch_next(idx, aiters[idx]))
+            pending.add(new_task)
+
+
+async def _fetch_next[T](idx: int, ait: AsyncIterator[T]) -> tuple[int, T | None, bool]:
+    try:
+        value = await ait.__anext__()
+        return (idx, value, False)
+    except StopAsyncIteration:
+        return (idx, None, True)

streamish/ops/filter.py

@@ -0,0 +1,235 @@
+"""Filter operations."""
+
+from collections.abc import (
+    AsyncIterable,
+    AsyncIterator,
+    Callable,
+    Hashable,
+    Iterable,
+    Iterator,
+)
+from typing import overload
+
+from streamish._util import is_async_iterable
+
+__all__ = ["take", "skip", "take_while", "skip_while", "distinct", "distinct_by"]
+
+
+@overload
+def take[T](n: int, it: Iterable[T]) -> Iterator[T]: ...
+
+
+@overload
+def take[T](n: int, it: AsyncIterable[T]) -> AsyncIterator[T]: ...
+
+
+def take[T](
+    n: int, it: Iterable[T] | AsyncIterable[T]
+) -> Iterator[T] | AsyncIterator[T]:
+    """Take first n elements."""
+    if is_async_iterable(it):
+        return _take_async(n, it)  # type: ignore[arg-type]
+    return _take_sync(n, it)  # type: ignore[arg-type, return-value]
+
+
+def _take_sync[T](n: int, it: Iterable[T]) -> Iterator[T]:
+    count = 0
+    for item in it:
+        if count >= n:
+            break
+        yield item
+        count += 1
+
+
+async def _take_async[T](n: int, it: AsyncIterable[T]) -> AsyncIterator[T]:
+    count = 0
+    async for item in it:
+        if count >= n:
+            break
+        yield item
+        count += 1
+
+
+@overload
+def skip[T](n: int, it: Iterable[T]) -> Iterator[T]: ...
+
+
+@overload
+def skip[T](n: int, it: AsyncIterable[T]) -> AsyncIterator[T]: ...
+
+
+def skip[T](
+    n: int, it: Iterable[T] | AsyncIterable[T]
+) -> Iterator[T] | AsyncIterator[T]:
+    """Skip first n elements."""
+    if is_async_iterable(it):
+        return _skip_async(n, it)  # type: ignore[arg-type]
+    return _skip_sync(n, it)  # type: ignore[arg-type, return-value]
+
+
+def _skip_sync[T](n: int, it: Iterable[T]) -> Iterator[T]:
+    count = 0
+    for item in it:
+        if count < n:
+            count += 1
+            continue
+        yield item
+
+
+async def _skip_async[T](n: int, it: AsyncIterable[T]) -> AsyncIterator[T]:
+    count = 0
+    async for item in it:
+        if count < n:
+            count += 1
+            continue
+        yield item
+
+
+@overload
+def take_while[T](pred: Callable[[T], bool], it: Iterable[T]) -> Iterator[T]: ...
+
+
+@overload
+def take_while[T](
+    pred: Callable[[T], bool], it: AsyncIterable[T]
+) -> AsyncIterator[T]: ...
+
+
+def take_while[T](
+    pred: Callable[[T], bool], it: Iterable[T] | AsyncIterable[T]
+) -> Iterator[T] | AsyncIterator[T]:
+    """Take elements while predicate is true."""
+    if is_async_iterable(it):
+        return _take_while_async(pred, it)  # type: ignore[arg-type]
+    return _take_while_sync(pred, it)  # type: ignore[arg-type, return-value]
+
+
+def _take_while_sync[T](pred: Callable[[T], bool], it: Iterable[T]) -> Iterator[T]:
+    for item in it:
+        if not pred(item):
+            break
+        yield item
+
+
+async def _take_while_async[T](
+    pred: Callable[[T], bool], it: AsyncIterable[T]
+) -> AsyncIterator[T]:
+    async for item in it:
+        if not pred(item):
+            break
+        yield item
+
+
+@overload
+def skip_while[T](pred: Callable[[T], bool], it: Iterable[T]) -> Iterator[T]: ...
+
+
+@overload
+def skip_while[T](
+    pred: Callable[[T], bool], it: AsyncIterable[T]
+) -> AsyncIterator[T]: ...
+
+
+def skip_while[T](
+    pred: Callable[[T], bool], it: Iterable[T] | AsyncIterable[T]
+) -> Iterator[T] | AsyncIterator[T]:
+    """Skip elements while predicate is true."""
+    if is_async_iterable(it):
+        return _skip_while_async(pred, it)  # type: ignore[arg-type]
+    return _skip_while_sync(pred, it)  # type: ignore[arg-type, return-value]
+
+
+def _skip_while_sync[T](pred: Callable[[T], bool], it: Iterable[T]) -> Iterator[T]:
+    skipping = True
+    for item in it:
+        if skipping and pred(item):
+            continue
+        skipping = False
+        yield item
+
+
+async def _skip_while_async[T](
+    pred: Callable[[T], bool], it: AsyncIterable[T]
+) -> AsyncIterator[T]:
+    skipping = True
+    async for item in it:
+        if skipping and pred(item):
+            continue
+        skipping = False
+        yield item
+
+
+@overload
+def distinct[T: Hashable](it: Iterable[T]) -> Iterator[T]: ...
+
+
+@overload
+def distinct[T: Hashable](it: AsyncIterable[T]) -> AsyncIterator[T]: ...
+
+
+def distinct[T: Hashable](
+    it: Iterable[T] | AsyncIterable[T],
+) -> Iterator[T] | AsyncIterator[T]:
+    """Remove duplicates."""
+    if is_async_iterable(it):
+        return _distinct_async(it)  # type: ignore[arg-type]
+    return _distinct_sync(it)  # type: ignore[arg-type, return-value]
+
+
+def _distinct_sync[T: Hashable](it: Iterable[T]) -> Iterator[T]:
+    seen: set[T] = set()
+    for item in it:
+        if item not in seen:
+            seen.add(item)
+            yield item
+
+
+async def _distinct_async[T: Hashable](it: AsyncIterable[T]) -> AsyncIterator[T]:
+    seen: set[T] = set()
+    async for item in it:
+        if item not in seen:
+            seen.add(item)
+            yield item
+
+
+@overload
+def distinct_by[T, K: Hashable](
+    key_fn: Callable[[T], K], it: Iterable[T]
+) -> Iterator[T]: ...
+
+
+@overload
+def distinct_by[T, K: Hashable](
+    key_fn: Callable[[T], K], it: AsyncIterable[T]
+) -> AsyncIterator[T]: ...
+
+
+def distinct_by[T, K: Hashable](
+    key_fn: Callable[[T], K], it: Iterable[T] | AsyncIterable[T]
+) -> Iterator[T] | AsyncIterator[T]:
+    """Remove duplicates by key function."""
+    if is_async_iterable(it):
+        return _distinct_by_async(key_fn, it)  # type: ignore[arg-type]
+    return _distinct_by_sync(key_fn, it)  # type: ignore[arg-type, return-value]
+
+
+def _distinct_by_sync[T, K: Hashable](
+    key_fn: Callable[[T], K], it: Iterable[T]
+) -> Iterator[T]:
+    seen: set[K] = set()
+    for item in it:
+        key = key_fn(item)
+        if key not in seen:
+            seen.add(key)
+            yield item
+
+
+async def _distinct_by_async[T, K: Hashable](
+    key_fn: Callable[[T], K], it: AsyncIterable[T]
+) -> AsyncIterator[T]:
+    seen: set[K] = set()
+    async for item in it:
+        key = key_fn(item)
+        if key not in seen:
+            seen.add(key)
+            yield item

streamish/ops/group.py

@@ -0,0 +1,183 @@
+"""Group operations."""
+
+import asyncio
+from collections import deque
+from collections.abc import AsyncIterable, AsyncIterator, Callable, Iterable, Iterator
+from typing import overload
+
+from streamish._util import ensure_async_iterator, is_async_iterable
+
+__all__ = ["batch", "window", "partition", "partition_async"]
+
+
+@overload
+def batch[T](
+    size: int, it: Iterable[T], *, timeout: float | None = None
+) -> Iterator[list[T]]: ...
+
+
+@overload
+def batch[T](
+    size: int, it: AsyncIterable[T], *, timeout: float | None = None
+) -> AsyncIterator[list[T]]: ...
+
+
+def batch[T](
+    size: int,
+    it: Iterable[T] | AsyncIterable[T],
+    *,
+    timeout: float | None = None,
+) -> Iterator[list[T]] | AsyncIterator[list[T]]:
+    """Group elements into batches by size or timeout."""
+    if size <= 0:
+        raise ValueError("size must be positive")
+    if is_async_iterable(it) or timeout is not None:
+        return _batch_async(size, it, timeout)  # type: ignore[arg-type]
+    return _batch_sync(size, it)  # type: ignore[arg-type, return-value]
+
+
+def _batch_sync[T](size: int, it: Iterable[T]) -> Iterator[list[T]]:
+    current: list[T] = []
+    for item in it:
+        current.append(item)
+        if len(current) >= size:
+            yield current
+            current = []
+    if current:
+        yield current
+
+
+async def _batch_async[T](
+    size: int,
+    it: Iterable[T] | AsyncIterable[T],
+    timeout: float | None,
+) -> AsyncIterator[list[T]]:
+    current: list[T] = []
+
+    ait: AsyncIterator[T]
+    if is_async_iterable(it):
+        ait = it.__aiter__()  # type: ignore[union-attr]
+    else:
+        sync_it: Iterable[T] = it  # type: ignore[assignment]
+        ait = ensure_async_iterator(iter(sync_it))
+
+    if timeout is None:
+        async for item in ait:
+            current.append(item)
+            if len(current) >= size:
+                yield current
+                current = []
+        if current:
+            yield current
+    else:
+        # Use a task-based approach to avoid cancelling the iterator
+        async def get_next() -> T:
+            return await ait.__anext__()
+
+        pending_task: asyncio.Task[T] | None = None
+        while True:
+            try:
+                if pending_task is None:
+                    pending_task = asyncio.create_task(get_next())
+                item = await asyncio.wait_for(
+                    asyncio.shield(pending_task), timeout=timeout
+                )
+                pending_task = None  # Task completed, clear it
+                current.append(item)
+                if len(current) >= size:
+                    yield current
+                    current = []
+            except TimeoutError:
+                if current:
+                    yield current
+                    current = []
+                # pending_task is still running, will be awaited next iteration
+            except StopAsyncIteration:
+                if current:
+                    yield current
+                break
+
+
+@overload
+def window[T](size: int, it: Iterable[T], *, step: int = 1) -> Iterator[list[T]]: ...
+
+
+@overload
+def window[T](
+    size: int, it: AsyncIterable[T], *, step: int = 1
+) -> AsyncIterator[list[T]]: ...
+
+
+def window[T](
+    size: int, it: Iterable[T] | AsyncIterable[T], *, step: int = 1
+) -> Iterator[list[T]] | AsyncIterator[list[T]]:
+    """Sliding window over elements."""
+    if size <= 0:
+        raise ValueError("size must be positive")
+    if step <= 0:
+        raise ValueError("step must be positive")
+    if is_async_iterable(it):
+        return _window_async(size, it, step)  # type: ignore[arg-type]
+    return _window_sync(size, it, step)  # type: ignore[arg-type, return-value]
+
+
+def _window_sync[T](size: int, it: Iterable[T], step: int) -> Iterator[list[T]]:
+    buf: deque[T] = deque(maxlen=size)
+    skip = 0
+    for item in it:
+        if skip > 0:
+            skip -= 1
+            buf.append(item)
+            continue
+        buf.append(item)
+        if len(buf) == size:
+            yield list(buf)
+            skip = step - 1
+            for _ in range(min(step, size)):
+                if buf:
+                    buf.popleft()
+
+
+async def _window_async[T](
+    size: int, it: AsyncIterable[T], step: int
+) -> AsyncIterator[list[T]]:
+    buf: deque[T] = deque(maxlen=size)
+    skip = 0
+    async for item in it:
+        if skip > 0:
+            skip -= 1
+            buf.append(item)
+            continue
+        buf.append(item)
+        if len(buf) == size:
+            yield list(buf)
+            skip = step - 1
+            for _ in range(min(step, size)):
+                if buf:
+                    buf.popleft()
+
+
+def partition[T](pred: Callable[[T], bool], it: Iterable[T]) -> tuple[list[T], list[T]]:
+    """Split into (matches, non_matches). Terminal operation."""
+    matches: list[T] = []
+    non_matches: list[T] = []
+    for item in it:
+        if pred(item):
+            matches.append(item)
+        else:
+            non_matches.append(item)
+    return matches, non_matches
+
+
+async def partition_async[T](
+    pred: Callable[[T], bool], it: AsyncIterable[T]
+) -> tuple[list[T], list[T]]:
+    """Split into (matches, non_matches). Terminal operation (async)."""
+    matches: list[T] = []
+    non_matches: list[T] = []
+    async for item in it:
+        if pred(item):
+            matches.append(item)
+        else:
+            non_matches.append(item)
+    return matches, non_matches