batch.py

from itertools import groupby, islice
from typing import Iterable, cast

from toolz.curried import curry, map, pluck


def batched(iterable, n):
    """Batch data into tuples of length n.

    >>> list(batched('ABCDEF', 3))
    [('A', 'B', 'C'), ('D', 'E', 'F')]
    """
    if n < 1:
        raise ValueError('n must be at least one')
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch
        

@curry
def split(key, iterable):
    """Batch data into tuples of unique elements."""
    curried = map(tuple, pluck(1, groupby(iterable, key=key)))
    for batch in cast(Iterable, curried):
        yield batch