partition.py

from typing import Any, Iterable, List


def partition(coll: Iterable[Any], size: int, use_padding=False) -> Iterable[Any]:
    """Lazily partition the iterable by making batches of size `size`.

    Arguments:
    * coll: iterable to partition
    * size: size of the desired partitions
    * use_padding: whether to use padding for the last batch in case it’s smaller
        than the desired size. The default behaviour is to use no padding.
    """
    batch: List[Any] = []
    for item in coll:
        if len(batch) == size:
            yield tuple(batch)
            batch = []
        batch.append(item)
    if batch:
        last_batch = batch if not use_padding else batch + [None] * (len(batch) - size)
        yield tuple(last_batch)
        return