Split up an iterable into batches I was writing a program that had a list of e-mail messages. It needed to group them into batches, with the condition that each batch must have at max 20 messages, and be maximally 50k in size (no individual message is bigger than 50k). Obviously itertools.groupby() is a candidate for this, but I really didn't want to keep around a key variable, so I decided to write my own batcher:

def batch(iterable, condition_func=(lambda x:True), max_in_batch=None):
    def Generator():
        nonlocal n, iterable, condition_func, max_in_batch
        num_in_batch = 1
        yield n
        while True:
            n = next(iterable)
            if num_in_batch == max_in_batch or condition_func(n):
                return
            num_in_batch += 1
            yield n

    n = next(iterable)
    while True:
        yield Generator()


Sample run:

&gt;&gt;&gt; gen = batch(itertools.count(1),(lambda x:x%10==9),7)
&gt;&gt;&gt; list(next(gen))
[1, 2, 3, 4, 5, 6, 7]
&gt;&gt;&gt; list(next(gen))
[8]
&gt;&gt;&gt; list(next(gen))
[9, 10, 11, 12, 13, 14, 15]
&gt;&gt;&gt; list(next(gen))
[16, 17, 18]
&gt;&gt;&gt; list(next(gen))
[19, 20, 21, 22, 23, 24, 25]
&gt;&gt;&gt; list(next(gen))
[26, 27, 28]


What do you think? Is this Pythonic?
Bugs

&gt;&gt;&gt; gen = batch(range(50), lambda x:x%10==9, 7)
&gt;&gt;&gt; for i in gen:
    print(list(i))


Traceback (most recent call last):
  File "&lt;pyshell#149&gt;", line 1, in &lt;module&gt;
    for i in gen:
  File "&lt;pyshell#142&gt;", line 13, in batch
    n = next(iterable)
TypeError: 'range' object is not an iterator


Damn, I expected to be able to use iterables as the iterable parameter, not iterators.

Let's fix it:

def batch(iterable, condition_func=(lambda x:True), max_in_batch=None):
    iterator = iter(iterable)
    def Generator():
        nonlocal n, iterator, condition_func, max_in_batch
        num_in_batch = 1
        yield n
        while True:
            n = next(iterator)
            if num_in_batch == max_in_batch or condition_func(n):
                return
            num_in_batch += 1
            yield n

    n = next(iterator)
    while True:
        yield Generator()


and try again:

&gt;&gt;&gt; gen = batch(range(50), lambda x:x%10==9, 7)
&gt;&gt;&gt; for i in gen:
    print(list(i))


[0, 1, 2, 3, 4, 5, 6]
[7, 8]
[9, 10, 11, 12, 13, 14, 15]
[16, 17, 18]
[19, 20, 21, 22, 23, 24, 25]
[26, 27, 28]
[29, 30, 31, 32, 33, 34, 35]
[36, 37, 38]
[39, 40, 41, 42, 43, 44, 45]
[46, 47, 48]
[49]
[49]
[49]
[49]
[49]
[49]
[49]
[49]
[49]
Traceback (most recent call last):
  File "&lt;pyshell#157&gt;", line 2, in &lt;module&gt;
    print(list(i))
  File "C:\Python33\lib\idlelib\PyShell.py", line 1336, in write
    return self.shell.write(s, self.tags)
KeyboardInterrupt


Dang… Let's try to limit that:

def batch(iterable, condition_func=(lambda x:True), max_in_batch=None):
    iterator = iter(iterable)
    def Generator():
        nonlocal n, on_going, iterator, condition_func, max_in_batch
        num_in_batch = 1
        yield n
        on_going = False
        while True:
            n = next(iterator)
            on_going = True
            if num_in_batch == max_in_batch or condition_func(n):
                return
            num_in_batch += 1
            yield n

    n = next(iterator)
    on_going = True
    while on_going:
        yield Generator()


Way better!

Use a for loop to iterate over an iterator

Iterating over elements and maintaining a counter of how much elements were iterated over is usually done using a for loop combined with enumerate:

def Generator():
    nonlocal n, on_going, iterator, condition_func, max_in_batch
    yield n
    on_going = False
    # Start enumerate at 1 because we already yielded n
    for num_in_batch, item in enumerate(iterator, 1):
        on_going = True
        n = item
        if num_in_batch == max_in_batch or condition_func(n):
            break
        yield item


Using this construct also help handling the on_going flag (I don't like the way we have to change it twice to continue iterating): when the break is not reached, we know we have exhausted the iterator.

def Generator():
    nonlocal n, on_going, iterator, condition_func, max_in_batch
    yield n
    # Start enumerate at 1 because we already yielded n
    for num_in_batch, item in enumerate(iterator, 1):
        n = item
        if num_in_batch == max_in_batch or condition_func(n):
            break
        yield item
    else:
        on_going = False


Use closures instead of nonlocals

Specifying iterator, condition_func, and max_in_batch in the nonlocal clause is not necessary as they are not modified within Generator. Instead you should use standard scope resolution to do the job.

In fact, your use of nonlocal to get non-trivial side effects on n had me wondered if an inner function is really the way to go. Let's try OOP instead:

class Batch:
    def __init__(self, iterable, condition=(lambda x:True), limit=None):
        self.iterator = iter(iterable)
        self.condition = condition
        self.limit = limit
        try:
            self.current = next(self.iterator)
        except StopIteration:
            self.on_going = False
        else:
            self.on_going = True

    def group(self):
        yield self.current
        # start enumerate at 1 because we already yielded the last saved item
        for num, item in enumerate(self.iterator, 1):
            self.current = item
            if num == self.limit or self.condition(item):
                break
            yield item
        else:
            self.on_going = False

    def __iter__(self):
        while self.on_going:
            yield self.group()


A bit more verbose but it seems to be a bit more explicit too. You use it as you would for your batch function:

&gt;&gt;&gt; gen = Batch(range(50), lambda x:x%10==9, 7)
&gt;&gt;&gt; for i in gen:
    print(list(i))


[0, 1, 2, 3, 4, 5, 6]
[7, 8]
[9, 10, 11, 12, 13, 14, 15]
[16, 17, 18]
[19, 20, 21, 22, 23, 24, 25]
[26, 27, 28]
[29, 30, 31, 32, 33, 34, 35]
[36, 37, 38]
[39, 40, 41, 42, 43, 44, 45]
[46, 47, 48]
[49]