k_sized_groups.py

#!/usr/bin/env python

from itertools import groupby

def k_sized_groups(sequence, k):
    """
    Generates iterable groups with a maximum of size of K members from the
    given sequence.

    >>> [list(g) for g in k_sized_groups([1, 2, 3, 4], 2)]
    [[1, 2], [3, 4]]
    >>> [tuple(g) for g in k_sized_groups([1, 2, 3, 4], 3)]
    [(1, 2, 3), (4,)]
    """


    # HACK! nonlocal doesn't exist
    non_local = { 'i': 0, 'val': True }

    # This is also a HACK for it to work properly with groupby.
    def factor_of_k(_ignored):
        if non_local['i'] % k == 0:
            # Flip the value
            non_local['val'] = not non_local['val']
        non_local['i'] += 1
        return non_local['val']

    return (group for _val, group in groupby(sequence, factor_of_k))