2.11 Comprehensions, iterators and generators

What this is and why it exists

At some point you meet a file bigger than your machine's memory, and the loop that worked on the sample stops working on the real thing. Generators are the answer: they produce items one at a time instead of building a list, so the memory a pipeline uses stops depending on how much data goes through it. Comprehensions are the other half of this topic — the idiomatic way to build a collection when you do want one.

The vocabulary

• Iterable — the glossary's phrase: "an object capable of returning its members one at a time".
• Iterator — "an object representing a stream of data", which produces the next item on request and raises a stop signal when it runs out.
• Comprehension — an expression that builds a list, dict or set from an iterable in one line.
• Generator function — a function containing yield, which returns a generator when called.
• Generator expression — a comprehension in round brackets, which produces items lazily.
• Lazy — computing each item only when it is asked for.
• Pipeline — a chain of generators, each taking the previous one's items.

The mental model

A list is a room with everything already in it. A generator is a conveyor belt — one item arrives, you deal with it, and it is gone. That is the whole trade: constant memory in exchange for one forward pass, with no indexing and no length.

The glossary is precise about how the machinery works, and it is worth reading rather than skimming. Each yield "temporarily suspends processing, remembering the execution state (including local variables and pending try-statements). When the generator iterator resumes, it picks up where it left off (in contrast to functions which start fresh on every invocation)." So a generator function is not really a function that returns items — it is a function you can pause. That reframe makes generators stop being mysterious.

Comprehensions come first in practice because they are what you reach for most. A list comprehension builds a list; a dict or set comprehension builds those; and the same expression in round brackets builds a generator instead, which is the change that turns a memory problem into no problem at all. The rule for readability is a comprehension that fits on one line and contains at most one condition — beyond that, a for loop is clearer, and clarity is worth more than the idiom.

Then the trap, which produces wrong answers rather than errors and is therefore the worst kind. A generator is single-use. The glossary explains why: once an iterator is exhausted, further calls "raise StopIteration again", and asking an exhausted iterator for an iterator returns itself — "making it appear like an empty container". So computing a sum from a generator and then a count from the same generator gives you a correct sum and a count of zero, with nothing raised anywhere. If you need two passes, either materialise it into a list — accepting the memory — or build the generator twice.

The other trap is subtler and worth naming: laziness moves when things happen. Nothing inside a generator runs until somebody asks for an item, so an error in the first line of a generator function surfaces at the first iteration, possibly far from where you called it. And a generator reading a file must be consumed before the file is closed, which is why a generator that opens its own file should also be the thing that iterates it, inside the with.

itertools is the standard library's collection of generator tools, and four are worth memorising: one that chains several iterables into a single stream, one that slices an iterator the way you would slice a list, one that groups consecutive equal items, and one that produces combinations or products. They compose with your own generators, and reaching for them saves writing loops that are surprisingly fiddly to get right.

In code

Checked against the glossary and the itertools reference.

import json
from itertools import islice

# A comprehension builds the whole list; the same in round brackets does not.
squares = [n * n for n in range(10)]
lazy_squares = (n * n for n in range(10_000_000))

by_name = {r["name"]: r["score"] for r in rows}
subjects = {r["subject"] for r in rows}


def read_records(path):
    """Yield one parsed record at a time; memory does not grow with the file."""
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                yield json.loads(line)


def high_scores(records, threshold=80):
    for r in records:
        if r["score"] >= threshold:
            yield r


# A pipeline: nothing is read until the last line asks for an item.
records = read_records("scores.jsonl")
top = high_scores(records)
first_five = list(islice(top, 5))

Read the last four lines carefully, because they are the point. Creating the generators does no work at all. Only islice asking for five items causes the file to be opened, lines to be read, and records to be parsed — and only enough of them to satisfy the request. A file of ten million lines and a file of ten behave identically in memory.

And the trap, shown:

scores = (r["score"] for r in rows)
total = sum(scores)
count = sum(1 for _ in scores)    # 0 — the generator is already exhausted

What you should now be able to explain or do

Say what a generator trades for constant memory. Explain, in the glossary's terms, what yield does to a function's execution state. Convert a list comprehension into a generator expression and say what changed. Write a generator pipeline that processes a file larger than memory. Recognise the exhausted-generator bug from its symptom — a correct first answer and a zero second one — and give two fixes. Name four itertools tools and what each is for.

Check yourself