Comprehensions

TL;DR

A comprehension is a one-expression way to build a list, dict, set, or generator from an iterable. The shape is always [expr for x in iterable if condition]. They are not “fancy loops” — they are the primary idiom for building collections in Python. Reading code, you should treat a comprehension as one atomic step: “build the list of squares of evens.” If you find yourself reaching for a multi-line for loop with result.append(...) inside, try the comprehension first.

The four flavours:

[x*x for x in nums if x > 0]              # list
{x*x for x in nums if x > 0}              # set
{x: x*x for x in nums if x > 0}           # dict
(x*x for x in nums if x > 0)              # generator (lazy, no storage)

The line where comprehensions stop being clearer: when you nest more than two for clauses, when the body needs side effects (logging, exceptions, multi-line transforms), or when the condition is itself several expressions. At that point, switch to a real for loop. The goal is readable code, not the shortest code.

The picture in your head

A comprehension is “construct A from B by applying transformation T, keeping only items that satisfy filter F.” Every comprehension is just a desugaring of:

result = []
for x in iterable:
    if condition:
        result.append(expr)

into:

result = [expr for x in iterable if condition]

If the desugared loop is what you’d write anyway, the comprehension is the right form. If the desugared loop has any other statements in it — logging, error handling, partial updates of multiple structures — the comprehension is the wrong form.

The four flavours

List comprehension

>>> nums = [1, 2, 3, 4, 5]
>>> [x*x for x in nums]
[1, 4, 9, 16, 25]

>>> [x*x for x in nums if x % 2 == 0]
[4, 16]

Set comprehension

>>> sentence = "the quick brown fox jumps over the lazy dog"
>>> {len(w) for w in sentence.split()}
{3, 4, 5}                       # unique word lengths

Dict comprehension

>>> words = ["the", "quick", "brown", "fox"]
>>> {w: len(w) for w in words}
{'the': 3, 'quick': 5, 'brown': 5, 'fox': 3}

>>> # Invert a vocab dict (token -> id) into (id -> token)
>>> id_to_tok = {i: t for t, i in vocab.items()}

Generator expression

>>> gen = (x*x for x in range(10**8))     # zero memory cost up front
>>> sum(gen)                                # consumes the generator
333333328333333350000000

>>> # Equivalent list comprehension would allocate ~800 MB

The generator-expression flavour is the one beginners under-use. Any time you’re feeding a comprehension straight into sum(), max(), min(), any(), all(), or ''.join(...), drop the brackets — the generator avoids materialising the intermediate list.

total = sum(x*x for x in nums)            # not sum([x*x for x in nums])
joined = ", ".join(str(x) for x in nums)  # not join([str(x) for x in nums])

Filtering and transformation

The general shape:

[transform(x) for x in iterable if predicate(x)]

You can chain conditions with and, or, not:

clean = [t for t in tokens if t.isalpha() and len(t) >= 2 and t not in stopwords]

You can also use a conditional expression on the value side (which is not the same thing as a filter — it transforms every item, but differently):

labels = [1 if score > 0.5 else 0 for score in scores]

Mixing them is fine and common:

clipped = [x if -1 <= x <= 1 else 0 for x in values if x is not None]

— “for every non-None value, return it if in [-1, 1], otherwise 0.”

Nested comprehensions — once is fine, twice is too much

Two for clauses iterate as if they were nested loops in the same order:

>>> # Cartesian product
>>> [(i, j) for i in range(3) for j in range(3)]
[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]

>>> # Flatten a list of lists
>>> rows = [[1, 2], [3, 4], [5, 6]]
>>> [x for row in rows for x in row]
[1, 2, 3, 4, 5, 6]

Read order: outer-to-inner, left-to-right. The first for is the outermost loop. This is the same order you’d write the nested for loops in.

A nested list (a comprehension whose expression is itself a comprehension) is different — it builds a list of lists:

>>> [[i*j for j in range(3)] for i in range(3)]
[[0, 0, 0], [0, 1, 2], [0, 2, 4]]

Useful for building 2D tables. Fine. But three levels of nesting starts to become unreadable; switch to a for loop or a NumPy/Pandas vectorized operation.

A worked ML example — building a vocabulary

Take a tokenized corpus and build the standard (token -> id, id -> token) pair, dropping rare tokens.

from collections import Counter

corpus = [
    ["the", "cat", "sat", "on", "the", "mat"],
    ["the", "dog", "barked"],
    ["a", "cat", "and", "a", "dog"],
]

# Flatten and count
freqs = Counter(tok for sent in corpus for tok in sent)

# Drop tokens that appear fewer than 2 times
vocab_tokens = [tok for tok, n in freqs.most_common() if n >= 2]
# ['the', 'cat', 'a', 'dog']

# Build both directions
tok_to_id = {tok: i for i, tok in enumerate(vocab_tokens)}
id_to_tok = {i: tok for tok, i in tok_to_id.items()}

# Encode the corpus, keeping only known tokens
encoded = [[tok_to_id[t] for t in sent if t in tok_to_id] for sent in corpus]
# [[0, 1, 0], [0, 3], [2, 1, 2, 3]]

Five comprehensions, no explicit loops. Each one expresses one operation: count, filter, build forward, build reverse, encode. The flatten-then-count uses a generator expression inside Counter to avoid materialising the flat list.

When NOT to use a comprehension

• The transformation needs try / except. Comprehensions can’t catch exceptions per item. Use a loop, or wrap each item in a helper function that returns a sentinel on failure.
• You want to update multiple structures in one pass. “For each example, append to lengths and increment count.” Use a loop.
• Side effects. Comprehensions are for constructing values. If you’re calling print(x) inside a comprehension, you’ve stopped using it as a comprehension. Use a loop.
• The expression is more than ~80 chars long. Wrap in a helper function or use a loop. Readability beats cleverness.
• You need an else on the loop. Doesn’t exist for comprehensions. Use a real loop.

A common anti-pattern is the “comprehension just to call a function” trick:

[print(x) for x in items]   # don't do this

It builds a list of Nones, then throws it away. Just use a for loop. for x in items: print(x) is shorter and doesn’t lie about intent.

Generator expressions deserve their own section

A generator expression is a comprehension with () instead of []. It produces values lazily, one at a time, without ever building a list.

gen = (line.strip() for line in open("huge.txt"))   # never loads file into memory
for clean in gen:
    process(clean)

This pattern is how you process datasets that don’t fit in RAM. You chain transformations, each one a generator, and the data flows through without materialising any intermediate.

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

def parsed(lines):
    return (json.loads(l) for l in lines if l)

def labeled(records):
    return ((r["text"], r["label"]) for r in records if "label" in r)

# Pipeline. Nothing has happened yet — these are all lazy.
pipeline = labeled(parsed(lines("dataset.jsonl")))

# Now consume.
for text, label in pipeline:
    train_step(text, label)

Memory cost is one record at a time, regardless of file size. List comprehensions all the way through would OOM.

The catch: you can iterate a generator once. After it’s exhausted, it yields nothing. If you need two passes, build a list — or call the generator function twice.

Walrus operator in comprehensions

Python 3.8’s := (the walrus) lets you bind a value mid-expression. In comprehensions, this is occasionally useful for “compute once, use twice”:

# Bad — calls expensive_score(x) twice per item
selected = [expensive_score(x) for x in xs if expensive_score(x) > 0.5]

# Good — compute once, reuse
selected = [s for x in xs if (s := expensive_score(x)) > 0.5]

Useful, but easy to abuse. Reach for it only when “compute, filter, emit” without it would force a duplicate call.

Common gotchas

• Generator expressions are single-pass. Iterating twice yields nothing the second time. Build a list if you need to reuse.
• Variable scope leaks in Python 2 (not 3). In Python 3, the loop variable inside a comprehension does NOT leak into the enclosing scope. Don’t rely on it.
• {} is an empty dict, not an empty set. Use set().
• Nested-loop order is outer-to-inner. [x for row in rows for x in row], not [x for x in row for row in rows] (the latter is a NameError).
• Don’t comprehend over a comprehension you also want to filter on. Use the walrus operator, or a helper function, or just write a loop.

Where this shows up in ML

• Building vocabularies, label maps, ID maps. Dict comprehensions.
• Filtering examples by length, label, score. List comprehensions.
• Streaming datasets that don’t fit in memory. Generator expressions chained into PyTorch DataLoader via an IterableDataset.
• Stacking tensors from a per-example function. torch.stack([encode(x) for x in batch]).
• Aggregating metrics across folds. mean = sum(scores) / len(scores) from a generator expression.
• Constructing class weights. weights = {c: total / cnt for c, cnt in Counter(labels).items()}.

The defensive habit: if a for loop in your code only ever appends to one collection and has no other body, it should probably be a comprehension. If it does anything else, it should stay a loop.

Resources

• PEP 202 — List Comprehensions — peps.python.org — the original spec.
• PEP 274 — Dict Comprehensions — peps.python.org — added in 3.0.
• PEP 289 — Generator Expressions — peps.python.org — the lazy version.
• PEP 572 — Assignment Expressions (the walrus) — peps.python.org.
• Fluent Python (2nd ed.) — Chapter 2 — oreilly.com — comprehensions in the broader context of sequences.