Getting items in batches (Python recipe) by Hamish Lawson

ActiveState Code (http://code.activestate.com/recipes/303279/)

islice is just what I needed. Soon after submitting the recipe I began thinking of alternatives to building a list in memory for each batch. islice is just what I needed to do the basic batching, which I've incorporated in my revised version, but I keep the result of islice as an iterator rather than convert it to a list as I want to keep memory requirements to a minimum.

All batches must be exhausted immediately. You should add a warning that a batch has to be entirely consumed before you can proceed to the next one. Otherwise users may be puzzled by the following:

>>> len([bi for bi in batch(range(19), 3)])
19

>>> list([bi for bi in batch(range(19), 3)][3])
[3]

Catching the StopIteration is not necessary when you are essentially reraising it:

>>> def batch(iterable, size):
...     it = iter(iterable)
...     while True:
...             bi = islice(it, size)
...             yield chain([bi.next()], bi)
...
>>> map(list, batch(xrange(10), 3))
[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]

Another approach. The groupby() function is somewhat versatile if some imagination goes into defining the key= function.

from itertools import groupby

def batch(iterable, size):
    def ticker(x, s=size, a=[-1]):
        r = a[0] = a[0] + 1
        return r // s
    for k, g in groupby(iterable, ticker):
         yield g

Removed handling of StopIteration. Duh! Yes, catching and reraising StopIteration is redundant. Thanks! I've removed it.

What to do about batches not being consumed? Your comment about unconsumed batches got me to thinking about what the right behaviour should be. Compensate for a batch iterator not getting exhausted by internally consuming the source iterator? Or have each batch iterator begin where the last one finished, so producing an infinite number of batch iterators if none of them get exhausted?

Using a generator class instead. I just posted my own solution to the same batching problem; should have searched a bit more thoroughly first. Here is a generator class that does the same thing. For those (like me) not thoroughly conversant with itertools, it might be easier to understand:

class groupcount(object):
    """Accept a (possibly infinite) iterable and yield a succession
    of sub-iterators from it, each of which will yield N values.

    >>> gc = groupcount('abcdefghij', 3)
    >>> for subgroup in gc:
    ...     for item in subgroup:
    ...             print item,
    ...     print
    ...
    a b c
    d e f
    g h i
    j
    """

    def __init__(self, iterable, n=10):
        self.it = iter(iterable)
        self.n = n

    def __iter__(self):
        return self

    def next(self):
        self.ondeck = self.it.next()
        return self._groupby()

    def _groupby(self):
        yield self.ondeck
        for i in xrange(1, self.n):
            yield self.it.next()

You can also generate the key func from itertools:

from itertools import groupby, izip, count

def batch(iterable, size):
    c = count()
    for k, g in groupby(iterable, lambda x:c.next()//size):
         yield g

My suggestion for returning batches of non-fixed size:

def batches(iterable, size):
    source = iter(iterable)
    while True:
        chunk = [val for _, val in zip(range(size), source) if val is not None]
        if not chunk:
            raise StopIteration
        yield chunk

The following version is slightly simpler and users iterators better:

def batches(iterable, size):
    source = iter(iterable)
    while True:
        chunk = [val for _, val in izip(xrange(size), source)]
        if not chunk:
            raise StopIteration
        yield chunk