codekingpro/portable-devtools

1import itertools2from collections.abc import Iterable3from collections.abc import Iterator4from typing import TypeVar5 6T = TypeVar("T")7 8 9def window(10    iterator: Iterable[T], behind: int = 0, ahead: int = 011) -> Iterator[tuple[T | None, ...]]:12    """13    Sliding window for an iterator.14 15    Example:16        >>> for prev, i, nxt in window(range(10), 1, 1):17        >>>     print(prev, i, nxt)18 19        None 0 120        0 1 221        1 2 322        2 3 None23    """24    # TODO: move into utils25    iters: list[Iterator[T | None]] = list(itertools.tee(iterator, behind + 1 + ahead))26    for i in range(behind):27        iters[i] = itertools.chain((behind - i) * [None], iters[i])28    for i in range(ahead):29        iters[-1 - i] = itertools.islice(30            itertools.chain(iters[-1 - i], (ahead - i) * [None]), (ahead - i), None31        )32    return zip(*iters)33