Comprehensions

Why this matters

A huge share of everyday code is "make a new collection from an old one": square these numbers, keep only the valid rows, map names to lengths, collect the distinct domains in a list of email addresses. Python has a compact syntax for exactly this, the comprehension, borrowed from mathematical set-builder notation ({x² | x ∈ S, x odd}). Used well, comprehensions are shorter, clearer and somewhat faster than the equivalent loop. Used badly, they become unreadable puzzles, so this lesson covers both the syntax and its limits.

From loop to list comprehension

Here is the classic build-a-list loop:

squares = []
for n in range(6):
    squares.append(n * n)
print(squares)   # [0, 1, 4, 9, 16, 25]

The same as a list comprehension:

squares = [n * n for n in range(6)]

Read it left to right as "a list of n * n, for each n in range(6)". The pieces:

[ n * n    for n in range(6) ]
  ───┬──   ─────────┬───────
 expression     for clause
 (what to       (where values
  collect)       come from)

The expression can be anything: a method call, an f-string, a tuple, a function call.

names = ["ann", "bo"]
[name.title() for name in names]          # ['Ann', 'Bo']
[f"<{name}>" for name in names]           # ['<ann>', '<bo>']
[(name, len(name)) for name in names]     # [('ann', 3), ('bo', 2)]

Filtering with if

An if clause after the for keeps only the items for which the condition is true:

nums = [4, -1, 7, 0, -3, 2]
positives = [n for n in nums if n > 0]          # [4, 7, 2]
doubled_pos = [2 * n for n in nums if n > 0]    # [8, 14, 4]

The equivalent loop shows the order of operations:

doubled_pos = []
for n in nums:
    if n > 0:
        doubled_pos.append(2 * n)

Do not confuse the filter if (at the end, no else) with a conditional expression in the expression part, which transforms every item:

[n if n > 0 else 0 for n in nums]     # [4, 0, 7, 0, 0, 2]  clamp negatives to 0

Here every item produces an output; the x if cond else y decides which value.

Dict and set comprehensions

Swap the brackets for braces and you get the other two kinds.

A dict comprehension has key: value before the for:

words = ["apple", "fig", "banana"]
lengths = {w: len(w) for w in words}
# {'apple': 5, 'fig': 3, 'banana': 6}

prices = {"tea": 3.5, "cake": 4.0}
with_tax = {item: round(p * 1.25, 2) for item, p in prices.items()}
# {'tea': 4.38, 'cake': 5.0}

inverted = {v: k for k, v in {"a": 1, "b": 2}.items()}
# {1: 'a', 2: 'b'}

If two items produce the same key, the later one wins, just as with repeated assignment.

A set comprehension has a single expression in braces, and duplicates collapse:

emails = ["ann@x.org", "bo@y.com", "cy@x.org"]
domains = {e.split("@")[1] for e in emails}
# {'x.org', 'y.com'}  (in some order: sets are unordered)

Note that {} is an empty dict; an empty set is written set().

Nested for clauses

A comprehension may have several for clauses. They nest left to right, exactly like nested loops, the first being the outermost:

pairs = [(x, y) for x in [1, 2] for y in "ab"]
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]

is

pairs = []
for x in [1, 2]:
    for y in "ab":
        pairs.append((x, y))

A common use is flattening a list of lists:

grid = [[1, 2], [3, 4], [5]]
flat = [x for row in grid for x in row]   # [1, 2, 3, 4, 5]

The order "for row in grid for x in row" feels backwards to many people at first. Say it as nested loops: "for each row in grid, for each x in that row, collect x".

A comprehension inside a comprehension is different: it produces a list of lists:

table = [[r * c for c in range(1, 4)] for r in range(1, 4)]
# [[1, 2, 3], [2, 4, 6], [3, 6, 9]]

The inner comprehension runs once per value of the outer r, producing one row each time.

What happens underneath

A comprehension compiles to a tight loop that appends to a new object, using a dedicated LIST_APPEND bytecode instead of looking up and calling the .append method each time. That is why comprehensions are typically a bit faster than the explicit loop. Since Python 3.12 comprehensions are also inlined into the surrounding code rather than run as a hidden nested function, which removed the overhead of a function call.

The loop variable is still local to the comprehension: it does not leak out.

x = "outer"
squares = [x * x for x in range(3)]
print(x)   # outer

(In Python 2 it leaked; in Python 3 it never does.)

Generator expressions: a preview

Write a comprehension with parentheses and you get a generator expression. It does not build a list; it produces an iterator that computes each value on demand:

squares = (n * n for n in range(1_000_000))
print(next(squares), next(squares))   # 0 1

When a generator expression is the only argument to a function, you can drop the extra parentheses:

total = sum(n * n for n in range(1_000_000))

This sums a million squares without ever holding a million-element list in memory. Like every iterator, it is one-shot. Module 13 covers generators properly; for now, prefer this form when you are going to consume the values once, for example in sum, max, any, or "".join(...).

Readability limits

Comprehensions are for building a collection from an iterable. Some guidelines:

• One for and maybe one if: almost always clear. Two for clauses: fine for flattening or pairs. More than that, or long conditions: write the loop.
• Do not use a comprehension for its side effects: [print(x) for x in items] builds a useless list of Nones. Use a for loop.
• If the expression needs a name to be understood, give it a function: [parse_row(line) for line in lines] beats a 90-character expression.
• Break long comprehensions over lines inside the brackets:

active_emails = [
    user["email"].lower()
    for user in users
    if user["active"] and "@" in user["email"]
]

Common mistakes

Putting the filter if first or with an else.

[n for n in nums if n > 0 else 0]

SyntaxError: invalid syntax

A filter if takes no else. If you meant "replace", move it to the front: [n if n > 0 else 0 for n in nums].

Expecting {} to be a set. type({}) is dict. Use set() for an empty set.

Wrong nesting order. [x for x in row for row in grid] raises NameError: name 'row' is not defined (or worse, silently uses an old row variable), because the first for runs first. Write the clauses in the order of the equivalent nested loops.

Shared inner lists. [[0] * 3] * 3 makes three references to one inner list (the aliasing trap from the collections module). [[0] * 3 for _ in range(3)] builds three separate lists. _ is the conventional name for a loop variable you do not use.

Recap

• [expr for x in iterable if cond] builds a list; {k: v for ...} a dict; {expr for ...} a set.
• The trailing if filters; a if cond else b in the expression transforms.
• Multiple for clauses nest left to right, outermost first; a comprehension inside a comprehension builds nested lists.
• The loop variable stays inside the comprehension.
• (expr for x in it) is a lazy generator expression; ideal inside sum, max, any, join.
• When a comprehension gets hard to read, write a loop.

Going further: in the REPL, compare dis.dis("[x*x for x in xs]") with the bytecode of the equivalent loop and spot LIST_APPEND. Then read PEP 709 (inlined comprehensions) for the story of the 3.12 change.