Seven items in batches of three. The third batch is the whole question.
items = [1, 2, 3, 4, 5, 6, 7]
# range with a step walks the starting positions
for start in range(0, len(items), 3):
print(items[start:start + 3])
# -> [1, 2, 3]
# [4, 5, 6]
# [7] short, and a slice gives it to you without complaint
# Dropping it: stop before any start that cannot fill a batch
for start in range(0, len(items) - 3 + 1, 3):
print(items[start:start + 3])
# -> [1, 2, 3]
# [4, 5, 6]
- The slice handles the short batch for you
items[6:9] on a seven-item list gives [7] rather than raising, because a slice clamps — so the ragged end needs no arithmetic at all. That is why this version is three lines while the index-counting version has an off-by-one waiting in it.
- Dropping the last batch is a change to the range, not a check at the end
len(items) - size + 1 as the stop means the loop never starts a batch it cannot fill. Handling it the other way — yielding and then deciding whether to take it back — is not possible in a generator, because it has already left. Decide before you yield.
- Both answers are right, in different places
In training, drop_last=True is normal: a final batch of 3 when the rest are 256 produces a noisy gradient step and, with batch normalisation, can fail outright on a batch of one. When you are inserting rows or calling an API, dropping it means losing data, and keeping it is the only defensible option. The flag exists because the caller knows which they are doing.
One loop over the starting positions, yielding a slice each time. range takes a step, which is exactly the batch size.
range(0, len(items), size) gives the starts, and items[start:start + size] is the batch — the slice handles the short one on its own. For drop_last, stop the range earlier so a short batch is never started.
The stop is len(items) - size + 1 when dropping, and len(items) otherwise. Validate the size before the range, or range raises its own error about a step of zero. Because the function contains a yield, the validation only runs when the generator is first iterated — which is worth knowing and, here, harmless.
def batches(items, size, drop_last=False): if size < ____: raise ____(f’size must be at least 1, got {size}') stop = len(items) - size + 1 if ____ else len(items) for start in range(0, stop, ____): ____ items[start:start + size]
def batches(items, size, drop_last=False):
if size < 1:
raise ValueError(f"size must be at least 1, got {size}")
# Stopping earlier is how a short final batch is never started —
# a generator cannot un-yield one it has already handed out.
stop = len(items) - size + 1 if drop_last else len(items)
for start in range(0, stop, size):
# The slice clamps, so the last batch is short rather than an error.
yield items[start:start + size]
With seven items and a size of three, the starts are 0, 3 and 6 — and 6 cannot fill a batch. A stop of 7 − 3 + 1 = 5 excludes it and keeps 3, which is the last one that can. Writing len − size gives a stop of 4, which also excludes 6 here and wrongly excludes a valid start whenever the division is exact. The +1 is doing real work.
A def containing a yield returns a generator immediately and runs no code at all until it is iterated — so a bad size raises on the first next(), not at the call. That is surprising the first time you meet it, and it is why a library wanting the error up front splits the function in two: a plain wrapper that validates and returns an inner generator.
Each batch is a new list holding the caller’s items, so the input survives and can be batched again. The consuming version — while items: items.pop(0) — works once, destroys the data, and is quadratic into the bargain, because popping from the front of a list shifts everything after it.
itertools.batched does exactly this from Python 3.12, and yields tuples; before that the itertools recipes have a version, and more_itertools has chunked and ichunked. All of them work on any iterable rather than needing a sequence with a length, which is the real limitation here — this function cannot batch a stream, and a stream is usually why you were batching. For training, use your framework’s DataLoader: it has drop_last, shuffling, and workers filling batches in the background while the last one is still being used, which is most of the actual benefit.
- Make it work on any iterable using itertools.islice, and notice that you can no longer compute the stop in advance.
- Add a pad value so the short batch is filled out to size, and decide whether a padded batch should be distinguishable from a full one.
- Compare against itertools.batched on the same inputs, and find the one difference in what they yield.