Comprehensions
14 min
A list comprehension builds a new list from an iterable in a single
expression — it's the Pythonic replacement for a for loop that just
appends to a list.
squares = [n * n for n in range(5)]
print(squares) # [0, 1, 4, 9, 16]
That's equivalent to:
squares = []
for n in range(5):
squares.append(n * n)
Add an if clause at the end to filter which elements get included:
evens = [n for n in range(10) if n % 2 == 0]
print(evens) # [0, 2, 4, 6, 8]
The transform (w.upper()) and the filter (if len(w) <= 3) can combine freely — this reads as 'uppercase every word, but only the short ones.'
The same syntax works for dicts and sets, just with different brackets:
lengths = {w: len(w) for w in words} # dict comprehension: {}
unique_lengths = {len(w) for w in words} # set comprehension: {}
Use a comprehension when the logic is a simple transform-and-maybe-filter that fits on one line — reach for a regular loop instead once the logic needs multiple steps or branches, since a comprehension crammed with nested conditions gets hard to read fast.
Write a function `squares_of_evens(nums)` that returns a list of the squares of the even numbers in `nums`, using a single list comprehension.
Write a function `word_lengths(words)` that returns a dict mapping each word in `words` to its length, using a dict comprehension.
What does `[x for x in range(5) if x % 2 == 0]` evaluate to?
You can write list and dict comprehensions with an optional filter, and know when a plain loop is more readable instead.