omlish 0.0.0.dev164__py3-none-any.whl → 0.0.0.dev166__py3-none-any.whl

omlish/iterators.py

@@ -1,300 +0,0 @@
-import collections
-import dataclasses as dc
-import functools
-import heapq
-import itertools
-import typing as ta
-
-# from . import check
-from . import lang
-
-
-T = ta.TypeVar('T')
-U = ta.TypeVar('U')
-
-_MISSING = object()
-
-
-class PeekIterator(ta.Iterator[T]):
-
-    def __init__(self, it: ta.Iterable[T]) -> None:
-        super().__init__()
-
-        self._it = iter(it)
-        self._pos = -1
-        self._next_item: ta.Any = _MISSING
-
-    _item: T
-
-    def __iter__(self) -> ta.Self:
-        return self
-
-    @property
-    def done(self) -> bool:
-        try:
-            self.peek()
-        except StopIteration:
-            return True
-        else:
-            return False
-
-    def __next__(self) -> T:
-        if self._next_item is not _MISSING:
-            self._item = ta.cast(T, self._next_item)
-            self._next_item = _MISSING
-        else:
-            self._item = next(self._it)
-        self._pos += 1
-        return self._item
-
-    def peek(self) -> T:
-        if self._next_item is not _MISSING:
-            return ta.cast(T, self._next_item)
-        self._next_item = next(self._it)
-        return self._next_item
-
-    def next_peek(self) -> T:
-        next(self)
-        return self.peek()
-
-    def takewhile(self, fn: ta.Callable[[T], bool]) -> ta.Iterator[T]:
-        while fn(self.peek()):
-            yield next(self)
-
-    def skipwhile(self, fn: ta.Callable[[T], bool]) -> None:
-        while fn(self.peek()):
-            next(self)
-
-    def takeuntil(self, fn: ta.Callable[[T], bool]) -> ta.Iterator[T]:
-        return self.takewhile(lambda e: not fn(e))
-
-    def skipuntil(self, fn: ta.Callable[[T], bool]) -> None:
-        self.skipwhile(lambda e: not fn(e))
-
-    def takethrough(self, pos: int) -> ta.Iterator[T]:
-        return self.takewhile(lambda _: self._pos < pos)
-
-    def skipthrough(self, pos: int) -> None:
-        self.skipwhile(lambda _: self._pos < pos)
-
-    def taketo(self, pos: int) -> ta.Iterator[T]:
-        return self.takethrough(pos - 1)
-
-    def skipto(self, pos: int) -> None:
-        self.skipthrough(pos - 1)
-
-
-class ProxyIterator(ta.Iterator[T]):
-
-    def __init__(self, fn: ta.Callable[[], T]) -> None:
-        self._fn = fn
-
-    def __iter__(self) -> ta.Self:
-        return self
-
-    def __next__(self) -> T:
-        return self._fn()
-
-
-class PrefetchIterator(ta.Iterator[T]):
-
-    def __init__(self, fn: ta.Callable[[], T] | None = None) -> None:
-        super().__init__()
-
-        self._fn = fn
-        self._deque: collections.deque[T] = collections.deque()
-
-    def __iter__(self) -> ta.Self:
-        return self
-
-    def push(self, item) -> None:
-        self._deque.append(item)
-
-    def __next__(self) -> T:
-        try:
-            return self._deque.popleft()
-        except IndexError:
-            if self._fn is None:
-                raise StopIteration from None
-        return self._fn()
-
-
-class RetainIterator(ta.Iterator[T]):
-
-    def __init__(self, fn: ta.Callable[[], T]) -> None:
-        super().__init__()
-
-        self._fn = fn
-        self._deque: collections.deque[T] = collections.deque()
-
-    def __iter__(self) -> ta.Self:
-        return self
-
-    def pop(self) -> None:
-        self._deque.popleft()
-
-    def __next__(self) -> T:
-        item = self._fn()
-        self._deque.append(item)
-        return item
-
-
-def unzip(it: ta.Iterable[T], width: int | None = None) -> list:
-    if width is None:
-        if not isinstance(it, PeekIterator):
-            it = PeekIterator(iter(it))
-        try:
-            width = len(it.peek())
-        except StopIteration:
-            return []
-
-    its: list[PrefetchIterator[T]] = []
-    running = True
-
-    def next_fn(idx):
-        nonlocal running
-        if not running:
-            raise StopIteration
-        try:
-            items = next(it)  # type: ignore
-        except StopIteration:
-            running = False
-            raise
-        for item_idx, item in enumerate(items):
-            its[item_idx].push(item)
-        return next(its[idx])
-
-    its.extend(PrefetchIterator(functools.partial(next_fn, idx)) for idx in range(width))
-    return its
-
-
-def take(n: int, iterable: ta.Iterable[T]) -> list[T]:
-    return list(itertools.islice(iterable, n))
-
-
-def chunk(n: int, iterable: ta.Iterable[T], strict: bool = False) -> ta.Iterator[list[T]]:
-    iterator = iter(functools.partial(take, n, iter(iterable)), [])
-    if strict:
-        def ret():
-            for chunk in iterator:
-                if len(chunk) != n:
-                    raise ValueError('iterable is not divisible by n.')
-                yield chunk
-        return iter(ret())
-    else:
-        return iterator
-
-
-def merge_on(
-        function: ta.Callable[[T], U],
-        *its: ta.Iterable[T],
-) -> ta.Iterator[tuple[U, list[tuple[int, T]]]]:
-    indexed_its = [
-        (
-            (function(item), it_idx, item)
-            for it_idx, item in zip(itertools.repeat(it_idx), it)
-        )
-        for it_idx, it in enumerate(its)
-    ]
-
-    grouped_indexed_its = itertools.groupby(
-        heapq.merge(*indexed_its),
-        key=lambda item_tuple: item_tuple[0],
-    )
-
-    return (
-        (fn_item, [(it_idx, item) for _, it_idx, item in grp])
-        for fn_item, grp in grouped_indexed_its
-    )
-
-
-def expand_indexed_pairs(
-        seq: ta.Iterable[tuple[int, T]],
-        default: T,
-        *,
-        width: int | None = None,
-) -> list[T]:
-    width_ = width
-    if width_ is None:
-        width_ = (max(idx for idx, _ in seq) + 1) if seq else 0
-    result = [default] * width_
-    for idx, value in seq:
-        if idx < width_:
-            result[idx] = value
-    return result
-
-
-##
-# https://docs.python.org/3/library/itertools.html#itertools-recipes
-
-
-def sliding_window(it: ta.Iterable[T], n: int) -> ta.Iterator[tuple[T, ...]]:
-    # sliding_window('ABCDEFG', 4) -> ABCD BCDE CDEF DEFG
-    iterator = iter(it)
-    window = collections.deque(itertools.islice(iterator, n - 1), maxlen=n)
-    for x in iterator:
-        window.append(x)
-        yield tuple(window)
-
-
-##
-
-
-@dc.dataclass()
-class UniqueStats:
-    key: ta.Any
-    num_seen: int
-    first_idx: int
-    last_idx: int
-
-
-@dc.dataclass(frozen=True)
-class UniqueItem(ta.Generic[T]):
-    idx: int
-    item: T
-    stats: UniqueStats
-    out: lang.Maybe[T]
-
-
-class UniqueIterator(ta.Iterator[UniqueItem[T]]):
-    def __init__(
-            self,
-            it: ta.Iterable[T],
-            keyer: ta.Callable[[T], ta.Any] = lang.identity,
-    ) -> None:
-        super().__init__()
-        self._it = enumerate(it)
-        self._keyer = keyer
-
-        self.stats: dict[ta.Any, UniqueStats] = {}
-
-    def __next__(self) -> UniqueItem[T]:
-        idx, item = next(self._it)
-        key = self._keyer(item)
-
-        try:
-            stats = self.stats[key]
-
-        except KeyError:
-            stats = self.stats[key] = UniqueStats(
-                key,
-                num_seen=1,
-                first_idx=idx,
-                last_idx=idx,
-            )
-            return UniqueItem(
-                idx,
-                item,
-                stats,
-                lang.just(item),
-            )
-
-        else:
-            stats.num_seen += 1
-            stats.last_idx = idx
-            return UniqueItem(
-                idx,
-                item,
-                stats,
-                lang.empty(),
-            )