pyiter.windowed

 1from typing import Deque, Iterable, Iterator, List
 2from collections import deque
 3from .transform import Transform, T
 4
 5
 6class WindowedTransform(Transform[T, List[T]]):
 7    """
 8    A transform that yields windows of a given size from an iterable.
 9    If partial_windows is True, then windows that are smaller than the given size are yielded.
10    If partial_windows is False, then only windows that are exactly the given size are yielded.
11    """
12
13    def __init__(self, iter: Iterable[T], size: int, step: int, partial_windows: bool):
14        super().__init__(iter)
15        self.size = size
16        self.step = step
17        self.partial_windows = partial_windows
18
19    def __do_iter__(self) -> Iterator[List[T]]:
20        window: Deque[T] = deque(maxlen=self.size)
21        for e in self.iter:
22            window.append(e)
23            if len(window) == self.size:
24                yield list(window)
25            if len(window) == self.size:
26                for _ in range(self.step):
27                    window.popleft()
28
29        if self.partial_windows and len(window) > 0:
30            yield list(window)