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