pythonalgorithmgeneratorchunking

how to split an iterable in constant-size chunks


I am surprised I could not find a "batch" function that would take as input an iterable and return an iterable of iterables.

For example:

for i in batch(range(0,10), 1): print i
[0]
[1]
...
[9]

or:

for i in batch(range(0,10), 3): print i
[0,1,2]
[3,4,5]
[6,7,8]
[9]

Now, I wrote what I thought was a pretty simple generator:

def batch(iterable, n = 1):
   current_batch = []
   for item in iterable:
       current_batch.append(item)
       if len(current_batch) == n:
           yield current_batch
           current_batch = []
   if current_batch:
       yield current_batch

But the above does not give me what I would have expected:

for x in   batch(range(0,10),3): print x
[0]
[0, 1]
[0, 1, 2]
[3]
[3, 4]
[3, 4, 5]
[6]
[6, 7]
[6, 7, 8]
[9]

So, I have missed something and this probably shows my complete lack of understanding of python generators. Anyone would care to point me in the right direction ?

[Edit: I eventually realized that the above behavior happens only when I run this within ipython rather than python itself]


Solution

  • This is probably more efficient (faster)

    def batch(iterable, n=1):
        l = len(iterable)
        for ndx in range(0, l, n):
            yield iterable[ndx:min(ndx + n, l)]
    
    for x in batch(range(0, 10), 3):
        print x
    

    Example using list

    data = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] # list of data 
    
    for x in batch(data, 3):
        print(x)
    
    # Output
    
    [0, 1, 2]
    [3, 4, 5]
    [6, 7, 8]
    [9, 10]
    

    It avoids building new lists.