Iterators and Generators

TL;DR

An iterator is any object that produces values one at a time via next(it), raising StopIteration when done. A generator is the easiest way to write an iterator: a function that uses yield instead of return. Calling the function returns a generator object; you can then iterate over it.

Why this matters in ML: training data does not fit in memory. A 200 GB parquet file, a streamed Kafka topic, an infinite synthetic-data sampler — none of them can be a list. They have to be consumed incrementally. Generators are the language-level mechanism that makes this possible without ceremony.

The mental model: a generator is a function that pauses at each yield, hands control (and a value) back to the caller, and resumes from the same point when next called. State between yields lives in local variables. The function “runs” on demand, not all at once.

The picture in your head

A list is a finished cake on a plate — fully baked, takes its full size in the kitchen even if you only eat a slice. A generator is a recipe being read aloud one step at a time: “crack one egg” — pause, hand the egg to the cook, wait. “Crack another” — pause again. The recipe never holds all twelve eggs at once; it just produces them on demand.

The cost of a list is proportional to its length. The cost of a generator is proportional to the largest single value plus a tiny bit of frame state. For a million 1KB strings: list is ~1GB, generator is ~1KB at any moment.

The iterator protocol — what makes a thing iterable

There are exactly two methods.

class MyIterator:
    def __iter__(self):
        return self                # an iterator returns itself
    def __next__(self):
        # ... compute next value ...
        # raise StopIteration when exhausted
        return value

A iterable is anything that __iter__() can be called on to produce an iterator. list, tuple, set, dict, str, file objects, and anything you write that has __iter__. The for loop is sugar:

for x in items:
    body(x)

# is exactly:
it = iter(items)
while True:
    try:
        x = next(it)
    except StopIteration:
        break
    body(x)

You almost never write the protocol by hand. You write a generator function and Python writes the protocol for you.

Generator functions — yield

A function with yield in it is a generator function. Calling it does not run the body; it returns a generator object. Iterating that object runs the body up to the next yield, returns the yielded value, and pauses.

def count_up_to(n):
    print("starting")
    i = 0
    while i < n:
        yield i
        i += 1
    print("done")

>>> g = count_up_to(3)        # nothing printed yet
>>> next(g)
starting
0
>>> next(g)
1
>>> next(g)
2
>>> next(g)
done
Traceback (most recent call last):
  ...
StopIteration

Notice the timing: print("starting") runs on the first next(), not when the generator is constructed. print("done") runs after the loop exits, on the call that raises StopIteration.

Local variables (i here) are preserved between yields. That’s what makes generators useful for stateful streaming.

Why generators win on memory

Concrete numerical example. A function that produces all the squares from 1 to N.

import sys

def squares_list(n):
    return [i*i for i in range(n)]

def squares_gen(n):
    for i in range(n):
        yield i*i

>>> sys.getsizeof(squares_list(10**6))
8448728                 # ~8.4 MB
>>> sys.getsizeof(squares_gen(10**6))
208                     # 208 bytes — independent of n

The generator holds zero squares at any moment. It computes one when asked, then forgets it. The list holds all million squares simultaneously.

For ML, this is the difference between a dataloader that streams 200GB of parquet files versus one that crashes on MemoryError.

Generator expressions — comprehensions, but lazy

We covered these in PY 103. Recap:

gen = (x*x for x in range(10**8))   # 200 bytes
lst = [x*x for x in range(10**8)]   # ~800 MB, also takes minutes

The two are interchangeable except for memory and laziness. Use the generator expression by default; use the list comprehension only when you need indexing, multiple passes, or len().

Chaining generators — building pipelines

Generators compose. Each one transforms a stream and passes it on. The data flows through the pipeline one item at a time.

import json

def lines(path):
    with open(path) as f:
        for line in f:
            yield line.rstrip()

def parsed(line_stream):
    for line in line_stream:
        if line:
            yield json.loads(line)

def with_label(record_stream):
    for r in record_stream:
        if "label" in r:
            yield r["text"], r["label"]

def batched(pairs, batch_size):
    batch = []
    for pair in pairs:
        batch.append(pair)
        if len(batch) == batch_size:
            yield batch
            batch = []
    if batch:
        yield batch

# Wire the pipeline. None of this runs yet.
pipeline = batched(with_label(parsed(lines("dataset.jsonl"))), batch_size=32)

# Now consume. Each batch is built on demand; only one is in memory.
for batch in pipeline:
    train_step(batch)

This is the entire architecture of torch.utils.data.IterableDataset. You write generator functions; PyTorch consumes them in worker processes. No file ever loads in full.

itertools — the standard-library generator zoo

Every operation you’d want on iterators is in itertools. The ones worth memorising:

from itertools import islice, cycle, chain

# First 5 items of an infinite generator
first_five = list(islice(infinite_gen(), 5))

# Infinite epoch sampler — never raises StopIteration
forever = cycle(train_loader)

# Concatenate train and val for a single combined eval pass
combined = chain(train_loader, val_loader)

itertools is C-implemented and fast. Reach for it before writing your own.

Generators that consume — send(), yield from

Generators are technically two-way: value = yield x lets a generator receive a value sent in via gen.send(...). This is the substrate that async/await is built on, and the basis of “coroutines” in the old sense (pre-3.5).

In day-to-day code, the only yield-related feature you’ll reach for is yield from, which delegates iteration to a sub-iterator:

def all_records(paths):
    for p in paths:
        yield from records_in(p)   # equivalent to: for r in records_in(p): yield r

Cleaner and slightly faster than the explicit nested loop.

Common gotchas

• Single-pass. A generator is exhausted after one full iteration. g = (x*x for x in range(3)); list(g); list(g) gives [0, 1, 4], [].
• Iterating an exhausted generator silently yields nothing. No error. If your training loop “ran zero batches”, suspect this.
• len() doesn’t work on generators. Convert to a list (which defeats the purpose), or count by iterating, or use __length_hint__ if the source supports it.
• Side effects don’t happen until consumption. Logging in a generator function won’t run until somebody iterates. Surprising during debugging.
• Generators close on garbage collection. Don’t rely on it for resource cleanup. Use a try / finally block, or wrap in a context manager.
• itertools.tee materialises when consumers progress at very different rates. Don’t use it as a free way to “save” a generator — it just buffers the data.

When generators are the wrong tool

• You need random access (x[i]). Use a list.
• You’ll iterate the same data many times and computing it is cheap. Use a list.
• The values are tiny and the source is bounded and small. The laziness savings don’t matter; just use a comprehension.
• You need len(...) repeatedly. Lists know their size; generators don’t.

A worked ML example — an infinite balanced sampler

A common training pattern: yield batches forever, with each class appearing in equal proportion regardless of how skewed the dataset is.

import random
from collections import defaultdict
from itertools import cycle

def balanced_sampler(examples, batch_size):
    """Yield batches forever, with class balance per batch."""
    by_class = defaultdict(list)
    for x, y in examples:
        by_class[y].append(x)

    classes = list(by_class)
    per_class = batch_size // len(classes)

    # Cycle each class's pool — we never exhaust
    cyclers = {c: cycle(random.sample(by_class[c], len(by_class[c])))
               for c in classes}

    while True:
        batch = []
        for c in classes:
            for _ in range(per_class):
                batch.append((next(cyclers[c]), c))
        random.shuffle(batch)
        yield batch

Note the architecture: cycle makes each per-class iterator infinite, the outer while True makes the sampler infinite. Memory cost is exactly one batch plus the index pools. The sampler runs forever and the trainer decides when to stop.

This is the actual structure used in PyTorch’s WeightedRandomSampler and friends, just dressed up.

Where generators show up in real ML codebases

• torch.utils.data.IterableDataset — its __iter__ is your generator function. The standard pattern for streaming datasets.
• Hugging Face datasets library — dataset.iter(batch_size=...) returns a generator over batches, lazy.
• vllm / tgi token streaming — for token in model.stream(...) is a generator; each yielded token is a step from the inference engine.
• OpenAI / Anthropic SDKs streaming responses — for chunk in client.messages.stream(...) is a generator.
• Beautifulsoup / lxml iterparse — generator over XML / HTML elements to avoid loading the whole DOM.
• TensorFlow tf.data.Dataset — same pattern, different syntax.

The defensive habit: when you write a function that builds and returns a list, ask “is this list ever going to be huge, and is the consumer going to iterate it once?” If yes to both, change the return [...] to yield ... and skip the allocation.

Resources

• PEP 255 — Simple Generators — peps.python.org — where yield came from.
• PEP 380 — Delegating to a subgenerator (yield from) — peps.python.org.
• Python docs — itertools — docs.python.org — the canonical reference.
• David Beazley — Generator Tricks for Systems Programmers — dabeaz.com — the talk that taught a generation of Python engineers what generators are for.
• Fluent Python (2nd ed.) — Chapter 17 — oreilly.com — Iterators, Generators, and Classic Coroutines.