iterchunked.py

from itertools import chain, islice, tee
from more_itertools import consume


_marker = object()


class iterchunked:
    def __init__(self, iterable, n):
        self._source = iter(iterable)
        self._n = n
        self._init = False
    
    def __iter__(self):
        return self

    def __next__(self):
        if self._init:
            consume(self._source, self._n)
        else:
            self._init = True

        item = next(self._source, _marker)
        if item is _marker:
            raise StopIteration
        else:
            self._source = chain([item], self._source)

        self._source, it = tee(self._source, 2)
        return islice(it, self._n)