Chapter 16

Comprehensions

Comprehensions build one collection from another in a single expression (Control Flow introduced them). The idea originated in mathematical set-builder notation, and passed into functional programming. Haskell had list comprehensions, and Python borrowed them.

Comprehensions require a mental shift. With a loop you describe how to build the result: make an empty list, walk the input, test each item, and append the ones you want. With a comprehension you describe what the result is, as a single expression, and let Python build it. A comprehension is shorter. It reads like the definition of the result rather than a recipe for it, and one line replaces several lines of loop bookkeeping.

List Comprehensions

A list comprehension consists of:

An input sequence.
A variable representing members of the input sequence.
An optional predicate expression.
An output expression that builds one element of the output list from each member that satisfies the predicate.

Several examples in this chapter use the same input list:

# a_list.py
a_list = [1, "4", 9, "a", 0, 4]

The list comprehension selects integers from the list and squares them:

# list_comprehension.py
from a_list import a_list

squared_ints = [
    e ** 2 for e in a_list if isinstance(e, int)]
print(squared_ints)
#: [1, 81, 0, 16]

In this comprehension:

The iterator walks through each member e of the input sequence a_list.
The predicate checks if the member is an integer.
If the member is an integer, the output expression squares it and appends it to the output list.

The built-in functions map() and filter() with a lambda achieve the same results. filter() applies a predicate to a sequence and retains the members that pass it. It produces a lazy iterator, which list() expands into a list:

# filtering.py
from a_list import a_list

ints = list(filter(lambda e: isinstance(e, int), a_list))

if __name__ == "__main__":
    print(ints)
#: [1, 9, 0, 4]

map() applies a function to each member:

# mapping.py
from filtering import ints

print(list(map(lambda e: e ** 2, ints)))  # type: ignore
#: [1, 81, 0, 16]

The two combine into a single expression:

# map_and_filter.py
from a_list import a_list

print(list(map(lambda e: e ** 2,  # type: ignore
               filter(lambda e: isinstance(e, int),
                      a_list))))
#: [1, 81, 0, 16]

The map()/filter() form funnels every element through lambda calls, and is harder to read. The comprehension inlines the test and the expression, and its brackets show at a glance that it produces a list. map() and filter() pay off when the function already exists, map(str.strip, lines) rather than [line.strip() for line in lines]. So the lambda makes map_and_filter.py worse, not map(). Functional Foundations returns to the choice.

The # type: ignore comments mark a cost beyond readability. filter() with a lambda predicate does not narrow the element type, so ty still sees int | str coming out and rejects e ** 2. Pyright infers the mixed list literal as list[Unknown] and checks nothing there. The comprehension’s if isinstance(e, int) does narrow, so list_comprehension.py needs no such comment. filter() can narrow, but only when its predicate is a named function annotated to return TypeIs[int] or TypeGuard[int] rather than bool (the narrowing summary covers the pair). filter(None, items) is the other narrowing form. It drops the falsy values, and the type checker knows no None survives.

A comprehension has a scope of its own:

# comprehension_scope.py
e = "outer"
squares = [e ** 2 for e in range(4)]
print(squares, e)
#: [0, 1, 4, 9] outer
total = 0
running = [(total := total + n) for n in range(5)]
print(running, total)
#: [0, 1, 3, 6, 10] 10

A comprehension’s loop variable belongs to the comprehension. The e inside the brackets is a different name from the e outside them, so the outer e survives untouched. A for loop behaves the opposite way: its loop variable stays behind in the enclosing scope after the loop ends.

The walrus operator is the exception to that scope. total := total + n assigns in the enclosing scope, so total holds the running sum after the comprehension finishes. That leak is deliberate: it lets a comprehension accumulate a value without a separate loop. Two uses are a SyntaxError: a walrus that rebinds the comprehension’s own iteration variable, and a walrus in a comprehension inside a class body.

The running sum is a rare use of the walrus. The common one avoids computing the same value twice, once to filter and once to produce the output:

# walrus_filter.py
def cube_if_even(n: int) -> int | None:
    return n ** 3 if n % 2 == 0 else None

data = range(6)
cubes = [
    y for x in data if (y := cube_if_even(x)) is not None
]
print(cubes)
#: [0, 8, 64]

(y := cube_if_even(x)) calls cube_if_even once, binds its result to y, and the if tests that same result. The output expression then reuses y. Without the walrus, the filter and the output would each need their own call, cube_if_even(x) is not None and cube_if_even(x), computing it twice for every element that passes.

Set Comprehensions

Set comprehensions use the same principles as list comprehensions, with {} instead of []. Braces build a set when the comprehension produces one value per element, and a dict when it produces a key: value pair. The colon decides which. Python has no empty-set literal, since {} is an empty dict. Write set().

The following set comprehension normalizes each name (capital first letter, the rest lower case), keeps the names longer than one character, and collapses the duplicates and case variants:

# set_comprehension.py
names = ["Bob", "JOHN", "alice", "bob", "ALICE", "J", "Bob"]

unique = {name[0].upper() + name[1:].lower()
          for name in names if len(name) > 1}

print(sorted(unique))  # Sorted for stable display
#: ['Alice', 'Bob', 'John']

same = set([name[0].upper() + name[1:].lower()
            for name in names if len(name) > 1])

print(unique == same)
#: True

same builds a list with a list comprehension, then passes it to set(). The result is the same, but the throwaway list costs time and memory that the set comprehension avoids.

Dictionary Comprehensions

A dictionary comprehension builds a dict. Each element produces a key and a value, with an optional filter. Here each name becomes an upper-case key mapped to its length, keeping only the names longer than three characters:

# dict_comprehension.py
names = ["Arthur", "Lancelot", "Bedevere", "Ni", "Robin"]

lengths = {name.upper(): len(name)
           for name in names if len(name) > 3}
print(lengths)
#: {'ARTHUR': 6, 'LANCELOT': 8, 'BEDEVERE': 8, 'ROBIN': 5}

The three parts mirror the list comprehension: the for clause supplies each name, the if clause drops "Ni", and the key: value expression before for produces each entry.

A common variant swaps a dictionary’s keys and values to invert a lookup:

# invert_dict.py
seat_of = {"Arthur": 1, "Galahad": 2, "Robin": 1}

name_at = {seat: name for name, seat in seat_of.items()}
print(name_at)
#: {1: 'Robin', 2: 'Galahad'}

Inverting assumes the values are unique. Arthur and Robin both sit at seat 1. Robin, entered later, overwrites Arthur at key 1, the same rule any duplicate dictionary key follows. Arthur never appears in name_at.

Nested Comprehensions

An identity matrix of size n is an n by n square matrix with ones on the main diagonal and zeros elsewhere. Python represents such a matrix as a list of lists, where each sub-list is a row. The following comprehension generates an identity matrix:

# identity_matrix.py
from typing import Final

SIZE: Final[int] = 6

matrix = [[1 if col == row else 0 for col in range(SIZE)]
          for row in range(SIZE)]

for row in matrix:
    print(row)
#: [1, 0, 0, 0, 0, 0]
#: [0, 1, 0, 0, 0, 0]
#: [0, 0, 1, 0, 0, 0]
#: [0, 0, 0, 1, 0, 0]
#: [0, 0, 0, 0, 1, 0]
#: [0, 0, 0, 0, 0, 1]

Read a nested comprehension from the outside in, not left to right. The outer comprehension supplies row. For each row, the inner comprehension runs the full col loop and produces one sub-list. The inner for col sits to the left of the outer for row but runs inside it. The output expression sits first but runs last, once per innermost iteration.

1 if col == row else 0 is a conditional expression, not a filter. It sits in the output position, before the for, and decides what each element is. Every col still produces one. An if after the for, as in [e ** 2 for e in a_list if isinstance(e, int)], decides whether the comprehension produces an element at all. The positions are not interchangeable: [x for x in xs if a else b] is a SyntaxError. When you need both, write each in its own position: [x if a else b for x in xs if c].

Nesting one comprehension inside another builds a list of lists. Writing two for clauses in one comprehension flattens instead, producing a single list. Those clauses do read left to right, in the order the equivalent nested loops would appear:

# flatten.py
rows = [[1, 2], [3, 4], [5]]
print([x for row in rows for x in row])
#: [1, 2, 3, 4, 5]

If you get that order backward, row is not bound yet when Python evaluates the first for clause:

# flatten_wrong_order.py
rows = [[1, 2], [3, 4], [5]]
try:
    print([x for x in row for row in rows])  # type: ignore
except NameError as e:
    print(e)
#: name 'row' is not defined

Feeding the Iterator Clause

Everything to the right of in is an ordinary iterable expression, so anything that produces one works there.

Use zip() to walk two sequences together, taking one element from each:

# zip_pairs.py
names = ["a", "b", "c", "d"]
values = [1, 2, 3]
print([f"{n}={v}" for n, v in zip(names, values)])
#: ['a=1', 'b=2', 'c=3']

zip() stops at the end of the shorter sequence. Pass strict=True to make a length mismatch raise a ValueError instead of silently truncating.

Unpack a tuple in the for clause’s target, here a (name, function) pair applied to a value:

# zip_unpack.py
operations = [
    ("doubled", lambda v: v * 2),
    ("squared", lambda v: v ** 2),
]
values = [10, 3, 42]
print([
    f"{name}({v}) = {f(v)}"
    for (name, f), v in zip(operations, values)
])
#: ['doubled(10) = 20', 'squared(3) = 9']

values has a third element, and zip() drops it, as in zip_pairs.py.

Here’s a two-level list comprehension using Path.walk():

# path_walk_comprehension.py
import tempfile
from pathlib import Path

# Build a small tree to walk: two .py files and one to skip
with tempfile.TemporaryDirectory() as tmp:
    root = Path(tmp)
    (root / "pkg").mkdir()
    for name in ("main.py", "pkg/util.py", "pkg/notes.txt"):
        (root / name).write_text("")
    py_paths = [
        (dirpath / f).relative_to(root).as_posix()
        for dirpath, _, files in root.walk()
        for f in files if f.endswith(".py")
    ]

for path in sorted(py_paths):  # Sorted for stable output
    print(path)
#: main.py
#: pkg/util.py

tempfile.TemporaryDirectory() is a context manager that creates a scratch directory and deletes it, and everything in it, when the with block exits. That gives the example a throwaway file tree to walk, without touching any real files or leaving anything behind.

In the py_paths comprehension, the first for walks the directories and the second for walks the files in each, flattening the tree into one list of paths. The filter tests f.endswith(".py") on the bare filename rather than building a Path and reading its .suffix. That avoids constructing a Path for every file in the tree, including the ones the filter skips.

root.rglob("*.py") finds the same two files in one line, with no explicit walk and no comprehension at all. Try rglob() first: a glob pattern already says what you want. walk() earns its place when the filter needs more than a glob pattern can express, a file’s size or its contents rather than its name, say, or when the comprehension needs the directory structure itself, not just the files at the bottom of it.

A with block, unlike a function body, does not create a new scope. The assignment to py_paths sits inside the with, but the name is still visible afterward, in the for path in sorted(py_paths): line below it. The comprehension finishes building py_paths, as strings, while the directory still exists. The for loop runs after the directory disappears, and by then nothing needs the files. Turning those brackets into parentheses would break the program: a generator expression would not start walking until sorted() pulls on it, and that pull comes outside the with. Generator Expressions returns to that gap.

Breaking Up a Complex Comprehension

A comprehension earns its place when you can read it in one pass. You can nest more for and if clauses, or wrap the whole thing in another call, but each one you add makes the expression harder to read in one pass. Here, filtering, flattening, sorting, and formatting all run in a single expression:

# dense_comprehension.py
warehouses = {
    "East": [
        ("wrench", 12, 4.50),
        ("drill", 0, 9.00),
        ("hammer", 5, 2.25),
    ],
    "West": [
        ("wrench", 3, 4.75),
        ("sander", 8, 15.00),
    ],
}

report = [
    f"{wh}: {name} (${price:.2f})"
    for wh, name, price in sorted(
        [(wh, name, price)
         for wh, items in warehouses.items()
         for name, qty, price in items
         if qty > 0 and price < 10],
        key=lambda t: t[2])
]

if __name__ == "__main__":
    for line in report:
        print(line)
#: East: hammer ($2.25)
#: East: wrench ($4.50)
#: West: wrench ($4.75)

Reading this means untangling several questions at once: which items qualify, how the warehouses flatten together, in what order the result arrives, and how each line renders. A comprehension nested inside sorted(), itself nested inside the outer comprehension, does four jobs in one expression.

Split into named steps, the logic is all still there, and each step now states its own purpose:

# comprehension_steps.py
from dense_comprehension import warehouses

in_stock = [
    (wh, name, price)
    for wh, items in warehouses.items()
    for name, qty, price in items
    if qty > 0 and price < 10
]
in_stock.sort(key=lambda t: t[2])

report = [
    f"{wh}: {name} (${price:.2f})"
    for wh, name, price in in_stock
]

for line in report:
    print(line)
#: East: hammer ($2.25)
#: East: wrench ($4.50)
#: West: wrench ($4.75)

in_stock answers “which items qualify, flattened across warehouses.” sort() answers “in what order.” report answers “how each line renders.” Each name documents a stage of the pipeline, so a reader can follow the transformation one step at a time instead of parsing every step simultaneously. Use this split whenever a comprehension needs a comment to explain what it does.

Comprehensions Build, Loops Execute

A comprehension’s output expression can be any expression, including a call with a side effect, such as print(). A comprehension can run code for its side effect and throw away the list it builds:

# comprehension_side_effects.py
wasted = [print(n) for n in [1, 2, 3]]
print(wasted)
#: 1
#: 2
#: 3
#: [None, None, None]

The comprehension calls print() for its side effect. print() returns None, so wasted ends up holding three Nones, a list built and immediately discarded. Worse, a reader scanning [...] expects a meaningful collection, and this comprehension is a loop written with the wrong punctuation.

The idiomatic version says what it does:

# for_loop_side_effects.py
for n in [1, 2, 3]:
    print(n)
#: 1
#: 2
#: 3

The for loop prints the same values without building a wasted list. The brackets no longer suggest a collection the code never uses. Use a comprehension when you want the collection it produces, and a for loop when you want the side effect. If nothing assigns or uses a comprehension’s result, write it as a loop instead.

Generator Expressions

A comprehension evaluates eagerly, so it immediately builds the entire result in memory. For a large data set, that wastes time and space, especially if you consume the result only once. A generator expression uses the same syntax with parentheses instead of brackets, and produces its values one at a time, on demand:

# generator_expression.py
from itertools import islice

squares = (n ** 2 for n in range(1_000_000))
print(next(squares))
#: 0
print(next(squares))
#: 1
print(list(islice(squares, 3)))
#: [4, 9, 16]

No computation runs until you pull a value. next() produces them one at a time, and itertools.islice() takes a few without building the million-element list.

The parentheses do not make it a tuple comprehension. No such form exists. When you need a tuple, pass the generator expression to tuple().

A generator expression can also feed set() and dict():

# set_dict_from_genexp.py
words = ["pol", "parrot", "fjord", "ex"]

lengths = set(len(w) for w in words)
print(sorted(lengths))
#: [2, 3, 5, 6]

initials = dict((w, w[0]) for w in words)
print(initials)
#: {'pol': 'p', 'parrot': 'p', 'fjord': 'f', 'ex': 'e'}

A set or dict must hold every element, so no lazy set or dict exists: set(...) or dict(...) consumes the whole generator immediately. The set comprehension {len(w) for w in words} and the dict comprehension {w: w[0] for w in words} build the same results, read more directly, and are the better choice.

Use a generator expression when the consumer takes values one at a time and does not need them all at once, such as sum(), any(), all(), min(), or max():

# genexp_consumers.py
nums = range(1_000_000)

print(sum(n * n for n in nums))
#: 333332833333500000
print(any(n == 12_345 for n in nums))
#: True
print(max(len(str(n)) for n in nums))
#: 6

None of these builds an intermediate collection of a million items, and any() stops when it finds a match. str.join() is the exception: it needs two passes, one to size the result and one to fill it, so it converts its argument to a list first. A generator expression therefore saves nothing over a list comprehension there.

A generator expression needs no parentheses of its own when it is a function’s only argument. A second argument makes them required: sum(n * n for n in nums, 0) is a SyntaxError, and sum((n * n for n in nums), 0) is the fix.

genexp_consumers.py iterates nums three times because range is re-iterable: each for over it starts again at zero. A generator expression is not re-iterable:

# spent_generator.py
nums = (n for n in range(10))
print(sum(n * n for n in nums))
#: 285
print(any(n == 5 for n in nums))
#: False
print(list(nums))
#: []

It runs once, and after something consumes its values it is empty. sum() drains nums, so any() sees no elements and reports False instead of True, with no exception to say the question was never asked. When you must traverse something twice, either materialize it with list() or write the generator expression again.

A generator expression defers everything but one thing. Creating one evaluates the outermost iterable immediately:

# genexp_timing.py
def source() -> list[int]:
    print("source() called")
    return [1, 2, 3]

factor = 2
gen = (n * factor for n in source())
#: source() called
print("generator created")
#: generator created
factor = 10
print(list(gen))
#: [10, 20, 30]

source() runs as Python builds the generator expression, before the line below it prints. The output expression waits, so the code reads factor when list() pulls the values rather than at the generator’s creation. The answer is [10, 20, 30] instead of [2, 4, 6]. A list comprehension has no such gap: it reads everything at once. That gap is also why path_walk_comprehension.py uses brackets. Its outermost iterable, root.walk(), would be called at creation, but the walking and the filtering would wait for a consumer that arrives after the directory disappears. Iterators explores generators further, and Generators covers the values they receive as well as the ones they produce.

Unpacking in Comprehensions

path_walk_comprehension.py flattens a tree with two for clauses. Python 3.15 (PEP 798) adds a more direct way to flatten. The unpacking operators * and ** may appear in the output expression of a comprehension or generator expression, splicing each iterable or mapping into the result. PEP 798 extends the PEP 448 unpacking from [*a, *b] and {**d1, **d2} to the comprehension form, and replaces many uses of two-for comprehensions, itertools.chain(), and itertools.chain.from_iterable():

# unpacking_comprehensions.py
rows = [[1, 2], [3, 4], [5]]
dicts = [{"a": 1}, {"b": 2}, {"a": 3}]

# *
print([*row for row in rows])
#: [1, 2, 3, 4, 5]

# **
print({**d for d in dicts})
#: {'a': 3, 'b': 2}

# In a generator expression
flat = (*row for row in rows)
print(list(flat))
#: [1, 2, 3, 4, 5]

# Shallow: one level
print([*row for row in [[1, [2, 3]], [4]]])
#: [1, [2, 3], 4]
# Braces plus * build a set
print({*s for s in [{1, 2}, {3}]})
#: {1, 2, 3}

[*row for row in rows] reads as “splice each row in,” and produces the same flat list as the two-for [x for row in rows for x in row], while saying what it does more directly. It is a shallow flatten, splicing only the outer iterable, so the nested [2, 3] above comes through unflattened. ** does the same for dictionaries, merging each mapping with later keys winning. With braces, * builds a set and ** builds a dict. Everywhere else the colon decides between the two; here neither form has a colon, so the unpacking operator decides. The asynchronous generator form ((*a async for a in agen())) works the same way (Concurrency introduces async syntax).

Choosing a Form

The four forms are one expression with different delimiters, and that is why learning the list form teaches all four. Brackets when you want a list. Braces for a set, or for a dict when a colon separates a key from a value. Parentheses when the consumer takes values one at a time and does not need them all at once. A for loop when you want the side effect rather than the collection.

The delimiters also decide when the work runs. Every form but the parenthesized one runs to completion before the next statement, so you pay the cost of a comprehension where you wrote it. A generator expression defers that cost to whoever consumes it, and pays it only for the values the consumer pulls.

Exercises

Using a_list from a_list.py ([1, "4", 9, "a", 0, 4]), write a list comprehension that finds the string elements made only of digits (e.isdigit()), converts each to int with int(e), and squares it. The predicate must reject "a" so int() never sees it. Of the types in a_list, only str has isdigit(), so the predicate must test isinstance(e, str) before calling it.
In identity_matrix.py, change the comprehension to put 2 on the diagonal instead of 1, without adding a second pass over the result.
In dict_comprehension.py, add "Galahad" to names, then predict which entries the comprehension produces before running it, given the len(name) > 3 filter.
In set_comprehension.py, drop the if len(name) > 1 filter, and predict how many entries unique holds before running it. Explain why "J" does not collide with "JOHN".
comprehension_side_effects.py builds a list of Nones. Write a version that keeps the printing but produces a list the caller can use, then say whether a comprehension or a for loop is the right shape for it.
In unpacking_comprehensions.py, add a fourth entry {"a": 5, "c": 9} to dicts and predict what {**d for d in dicts} produces before running it, paying attention to which value wins for the key "a".
In spent_generator.py, move the any() line above the sum() line. Predict all three printed values before running it, remembering that any() stops when it finds a match.