summaryrefslogtreecommitdiffstats
path: root/utils.py
blob: d322ed01b501bed532890ce45daa58960e18ad29 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from collections import deque
from collections.abc import Iterable, Iterator, Generator
from itertools import islice
from typing import TypeVar

T = TypeVar('T')
def sliding_window(iterable: Iterable[T], n: int) -> Generator[tuple[T, ...], None, None]:
    # sliding_window('ABCDEFG', 4) -> ABCD BCDE CDEF DEFG
    it: Iterator[T] = iter(iterable)
    window: deque[T] = deque(islice(it, n), maxlen=n)
    if len(window) == n:
        yield tuple(window)
    for x in it:
        window.append(x)
        yield tuple(window)