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