1import itertools
2from typing import TypeVar, Iterable, Iterator, Tuple, Optional, List
3
4T = TypeVar('T')
5
6
7def window(iterator: Iterable[T], behind: int = 0, ahead: int = 0) -> Iterator[Tuple[Optional[T], ...]]:
8    """
9    Sliding window for an iterator.
10
11    Example:
12        >>> for prev, i, nxt in window(range(10), 1, 1):
13        >>>     print(prev, i, nxt)
14
15        None 0 1
16        0 1 2
17        1 2 3
18        2 3 None
19    """
20    # TODO: move into utils
21    iters: List[Iterator[Optional[T]]] = list(itertools.tee(iterator, behind + 1 + ahead))
22    for i in range(behind):
23        iters[i] = itertools.chain((behind - i) * [None], iters[i])
24    for i in range(ahead):
25        iters[-1 - i] = itertools.islice(
26            itertools.chain(iters[-1 - i], (ahead - i) * [None]),
27            (ahead - i),
28            None
29        )
30    return zip(*iters)
31