how to split an iterable in constant-size chunks
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.
FWIW, the recipes in the itertools module provides this example:
def grouper(n, iterable, fillvalue=None):
"grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
args = [iter(iterable)] * n
return zip_longest(fillvalue=fillvalue, *args)
It works like this:
>>> list(grouper(3, range(10)))
[(0, 1, 2), (3, 4, 5), (6, 7, 8), (9, None, None)]
As others have noted, the code you have given does exactly what you want. For another approach using itertools.islice
you could see an example of following recipe:
from itertools import islice, chain
def batch(iterable, size):
sourceiter = iter(iterable)
while True:
batchiter = islice(sourceiter, size)
yield chain([batchiter.next()], batchiter)
More-itertools includes two functions that do what you need:
-
chunked(iterable, n)
returns an iterable of lists, each of lengthn
(except the last one, which may be shorter); -
ichunked(iterable, n)
is similar, but returns an iterable of iterables instead.
Solution for Python 3.8 if you are working with iterables that don't define a len
function, and get exhausted:
from itertools import islice
def batcher(iterable, batch_size):
iterator = iter(iterable)
while batch := list(islice(iterator, batch_size)):
yield batch
Example usage:
def my_gen():
yield from range(10)
for batch in batcher(my_gen(), 3):
print(batch)
>>> [0, 1, 2]
>>> [3, 4, 5]
>>> [6, 7, 8]
>>> [9]
Could of course be implemented without the walrus operator as well.