Comprehensions
A loop says how to build a collection, one step at a time. A comprehension says what the collection is, and Python builds it.
Comprehensions build one collection from another in a single expression. 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. 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 filter expression.
- An output expression that builds one element of the output list from each member that passes the filter.
The first examples take their input from a list that mixes integers and strings:
# 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]
- The iterator walks through each member
eof the input sequencea_list. - The filter checks whether the member is an integer.
- If the member is an integer, the output expression squares it and the result joins the output list.
The map()
and filter() Equivalent
The built-in functions map() and
filter(), each given a lambda, produce
the same list that list_comprehension.py
builds. 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]Nesting the filter() call inside the
map() call does both steps in one 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.
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 # type: ignore comment.
filter() can narrow, but only when its predicate is
a named function annotated to return TypeIs[int]
or TypeGuard[int] rather than
bool. filter(None, items) is the other
narrowing form. It drops the falsy values, and the type checker
knows no None survives.
Scope and the Walrus Operator
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] 10A 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 scope works like a function’s, and that matters in a
class body. Names assigned in a class body are invisible to the
scopes nested inside it, so a comprehension there cannot read
the class’s other attributes. With base = 3 in the
class body, [n * base for n in range(3)] raises a
NameError. The outermost iterable is the one part
Python evaluates in the enclosing scope, so
[n for n in range(base)] works.
The walrus operator escapes the comprehension’s 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. Three
uses are a SyntaxError: a walrus that rebinds the
comprehension’s own iteration variable, a walrus inside the
iterable expression that follows in, 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 each need their own call,
cube_if_even(x) is not None and
cube_if_even(x), calling
cube_if_even() 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().
# 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)
#: Truesame builds a list with a list comprehension,
then passes it to set(). The result is identical,
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. To keep both names, group
instead of inverting: a for loop over
seat_of.items() that appends each name to a defaultdict(list)
keyed by seat produces
{1: ['Arthur', 'Robin'], 2: ['Galahad']}, a shape a
comprehension cannot accumulate.
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.
# 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.
The positions are not interchangeable.
[x for x in xs if a else b] is a
SyntaxError. When you need both a conditional
expression and a filter, 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 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 definedThe NameError depends on the enclosing scope
having no row. When an earlier statement binds one,
as the for row in matrix: loop in identity_matrix.py does,
the first clause iterates over that stale row, and
the comprehension produces a wrong list with no exception.
Supplying the Input Sequence
Everything to the right of in is an ordinary
iterable expression, so anything that produces one works
there.
zip() walks 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.
The for clause’s target can unpack a tuple, 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.
Path.walk() produces a
(dirpath, dirnames, filenames) tuple for each
directory in a tree, and a second for clause walks
the names in filenames:
# 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.pytempfile.TemporaryDirectory() is a context
manager that creates a scratch directory and deletes it, and
everything in it, when the with block exits. The
scratch directory gives the example a throwaway file tree to
walk, and leaves the file system as the example found it.
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. Testing the bare filename avoids
constructing a Path for every file in the tree,
including the ones the filter skips.
A with block, unlike a function body, runs in
the enclosing 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. If you turn the comprehension’s brackets into
parentheses, the program prints nothing and raises no exception.
A generator expression starts walking when sorted()
pulls on it, and that pull comes outside the with.
By then the context manager has deleted the directory, and
Path.walk() ignores the error unless you pass
on_error.
root.rglob("*.py") finds the same two files in
one line, with no explicit walk and no comprehension. Try
rglob() first. A glob pattern 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, not just the files at the bottom of it.
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 comprehension in another call, but
each one you add makes the expression harder to read.
# 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 the report comprehension 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. The outer comprehension
iterates over sorted() of an inner comprehension,
and that one expression does four jobs. Giving each stage a name
gives each question its own statement:
# 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]]
#: 1
#: 2
#: 3
print(wasted)
#: [None, None, None]The comprehension calls print() for its side
effect. print() returns None, so
wasted ends up holding three Nones, a
list that no code can use. Worse, a reader scanning
[...] expects a meaningful collection, and this
comprehension is a loop written with the wrong punctuation.
# for_loop_side_effects.py
for n in [1, 2, 3]:
print(n)
#: 1
#: 2
#: 3The for loop prints the same values without
building a wasted list. With the comprehension’s brackets gone,
nothing suggests an unused collection. 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]The generator computes a square each time 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 a generator expression 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
without a call to set() or dict(), 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))
#: 6None of these builds an intermediate collection of a million
items, and any() stops when it finds a match.
str.join() 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.
A Generator Expression Runs Once
genexp_consumers.py can
iterate over nums three times because a
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))
#: []A generator expression runs once, and after something
consumes its values it is empty. The generator keeps its place
between calls, and once its body has finished, every further
next() raises StopIteration (Iterators
has the protocol), which sum(), any(),
and list() each read as the end. sum()
drains nums, so any() sees no elements
and reports False instead of True,
with no exception to say the question went unasked. When you
must traverse something twice, either materialize it with
list() or write the generator expression again.
The Gap Between Creation and Consumption
A generator expression defers everything but one thing. Creating a generator expression evaluates its 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 generator created line
prints. The output expression waits, so it reads
factor when list() pulls the values
rather than when Python creates the generator. The answer is
[10, 20, 30] instead of [2, 4, 6]. A
list comprehension has no such gap. It reads everything at
once.
The gap between creation and consumption is also why path_walk_comprehension.py
uses brackets. With parentheses, Python calls
root.walk(), its outermost iterable, at creation,
but the walking and the filtering wait for a consumer that
arrives after the directory disappears. A lambda
created in a comprehension reads its variables late for the same
reason. Its body runs when something calls it. Function
Objects demonstrates that trap. 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
way to flatten with one for clause. The unpacking
operator * may appear in the output expression of a
comprehension or generator expression, and ** in a
dict comprehension, splicing each iterable or mapping into the
result. PEP 798 extends the PEP 448 unpacking
in displays, [*a, *b] and {**d1, **d2}
(Containers), to
the comprehension form. The new form 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], with one
for clause instead of two. 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,
so 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
This chapter’s solutions give a hint, usually the shape of the code, and a full answer for each exercise.
- Using
a_listfroma_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 tointwithint(e), and squares it. The filter must reject"a"beforeint()sees it. Of the types ina_list, onlystrhasisdigit(), so the filter must testisinstance(e, str)before calling it. - In
identity_matrix.py, change the comprehension to put2on the diagonal instead of1, without adding a second pass over the result. - In
dict_comprehension.py, add"Galahad"tonames, then predict which entries the comprehension produces before running it, given thelen(name) > 3filter. - In
set_comprehension.py, drop theif len(name) > 1filter, and predict how many entriesuniqueholds before running it. Explain why"J"does not collide with"JOHN". comprehension_side_effects.pybuilds a list ofNones. Write a version that keeps the printing but produces a list the caller can use, then say whether a comprehension or aforloop is the right shape for it.- In
unpacking_comprehensions.py, add a fourth entry{"a": 5, "c": 9}todictsand 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 theany()line above thesum()line. Predict all three printed values before running it, remembering thatany()stops when it finds a match. - In
genexp_timing.py, turn the generator expression into a list comprehension and name the resultbuilt. Predict the three printed lines, and their order, before running it. Explain which value offactorthe result uses.