6. Iteration and generators

Every for loop runs the same two-method protocol, and generators let you write that protocol as a function that pauses. This is how you stream datasets and LLM tokens without loading everything into memory.

The iterator protocol

• An iterable has __iter__(), which returns an iterator.
• An iterator has __next__(), which returns the next item or raises StopIteration when finished.

for x in xs: is roughly this:

it = iter(xs)
while True:
    try:
        x = next(it)
    except StopIteration:
        break
    ...  # loop body

Lists are iterables you can loop over again and again. Iterators are one-shot: once exhausted, they stay empty.

doubled = (x * 2 for x in range(3))
list(doubled)    # [0, 2, 4]
list(doubled)    # []  <- the classic "second epoch sees no data" bug

This is why PyTorch’s DataLoader is an iterable that creates a fresh iterator each epoch.

Generators

A function containing yield returns a generator object when called. Each next() runs the body until the next yield, then freezes all local state until the following next().

import json

def read_jsonl(path):
    with open(path, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                yield json.loads(line)

def batched(items, n):
    batch = []
    for item in items:
        batch.append(item)
        if len(batch) == n:
            yield batch
            batch = []
    if batch:
        yield batch

for batch in batched(read_jsonl("data.jsonl"), 32):
    embed(batch)

Memory use stays constant whatever the file size, because each stage pulls one item at a time from the stage before. That is a lazy pipeline. (Python 3.12 added itertools.batched, which does the same job.)

Laziness: benefits and traps

• Benefit: constant memory, processing can start before all data exists, and streams can be infinite.
• Trap: errors surface when the data is consumed, not where the pipeline was defined.
• Trap: a generator holding a file open keeps it open until it is exhausted or closed.
• Trap: len() does not work on a generator, because its length is unknown until it runs.

yield from

yield from sub() re-yields every item from another generator. It keeps recursive traversals short, such as walking nested JSON or directory trees.

Streaming LLM responses is the same pattern

Streaming APIs hand you an iterator of chunks (or an async iterator, Part 12) that yields each piece as the server sends it:

for event in stream:                       # arrives piece by piece
    if event.type == "content_block_delta":   # event names differ per SDK
        print(event.delta.text, end="", flush=True)

Event names and fields differ between SDKs, but the shape is universal: iterate over chunks as they arrive. Server-sent events (SSE) are the HTTP version of the same idea.

itertools: the toolbox

• islice(it, n): take the first n items of any iterator, for example to peek at a dataset.
• chain.from_iterable(lists): flatten one level of nesting.
• groupby(items, key): group consecutive items, so sort by the same key first.
• product, combinations, permutations: hyperparameter grids and pairings.
• accumulate: running totals.
• count, cycle, repeat: infinite streams.
• tee: split one iterator into several; it buffers, so watch memory.

Sending values into generators

gen.send(value) resumes a generator and passes a value in. That mechanism is the historical root of coroutines; modern code uses async def instead (Part 12). Know it exists, and expect to rarely write it.

Questions to ask:

• Is this an iterator that will be empty the second time through?
• Does this load everything into memory, or stream it?
• When does this generator actually run, and when does it release its resources?