Split List into Batches Using Generator in Python

Sometimes we need to split the list of items into batches of specified size. In this tutorial, we’ll discuss a method to split list into batches using a generator in Python.

Let’s consider we have a list of numbers. We will split this list into a small batch of size 3.

a = [1,2,3,4,5,6,7,8,10]

So the output would be like this:

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
[10]

Batch generator implementation

It is easy to create a generator function in Python. Generator functions uses yield statement instead of return statement. A simple generator function as shown below.

def gen():  # defines a generator function
    yield 42

The batch generator implementation has the following steps:

• Get the list and batch size.
• Get the start index for the batch using the range function.
• Slice the list using the start index and end index.
• Yield the result.

The generator function for the batch function can be defined as:

def generate_batch(lst, batch_size):
    """  Yields batch of specified size """
    for i in range(0, len(lst), batch_size):
        yield lst[i : i + batch_size]
       

Let’s examine this generator function line by line.

for i in range(0, len(lst), batch_size):

Here range function used to get the start index for the list slice.

 yield lst[i : i + batch_size]

Here list is sliced using the start index (i) and end index using ( i + batch_size).

test_list = list(range(1,11))

gen = generate_batch(test_list,3)

for batch in gen:
  print(batch)

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
[10]

For this function, batch_size must be greater than 0. Improved function as shown.

def generate_batch(lst, batch_size):
    """  Yields bacth of specified size """
    if batch_size<=0:
       return
    for i in range(0, len(lst), batch_size):
        yield lst[i : i + batch_size]
       

Conclusion

There you have it, a simple way to split the list into batches using a generator in Python. Read about list append extend performance.