Comprehensions — Python

TL;DR

[expression for item in iterable] builds a new list by applying the expression to every item.
An if at the end filters items; a if condition else b at the front chooses a value for every item.
Curly braces give a set comprehension, or a dictionary comprehension when the expression is key: value.
Parentheses give a generator expression, which produces its values one at a time instead of building a collection.
If a comprehension no longer fits comfortably on one or two lines, a plain for loop is usually clearer.

List comprehensions

A very common job is: take every item of a collection, do something with it, and collect the results in a new list. A list comprehension does that in a single expression. Its basic form is:

[expression for item in iterable]

Read it from the middle: for every item in the iterable, evaluate the expression, and put the result in the new list. The square brackets say that the result is a list.

celsius = [0, 10, 20, 30]
fahrenheit = [c * 9 / 5 + 32 for c in celsius]
print(fahrenheit)

names = ["ada", "linus", "grace"]
print([name.title() for name in names])
print([len(name) for name in names])

[32.0, 50.0, 68.0, 86.0]
['Ada', 'Linus', 'Grace']
[3, 5, 5]

The equivalent loop

Every list comprehension can be written as a for loop that starts with an empty list and appends to it. The two versions below build the same list; the comprehension just says it more directly, and runs slightly faster.

prices = [4, 12, 7, 20]

# with a loop
doubled = []
for price in prices:
    doubled.append(price * 2)
print(doubled)

# with a comprehension
doubled = [price * 2 for price in prices]
print(doubled)

[8, 24, 14, 40]
[8, 24, 14, 40]

The original list is not changed: a comprehension always builds a new list. It works on any iterable — a list, a tuple, a string, a range() or a file.

Filtering with if

Add an if clause at the end to keep only the items for which a condition is true. This combines two classic operations in one expression: map (transform each item) and filter (keep some items).

temperatures = [14, -2, 21, 0, 8, -5]

frost = [t for t in temperatures if t < 0]
print(frost)

squares_of_evens = [n * n for n in range(10) if n % 2 == 0]
print(squares_of_evens)

[-2, -5]
[0, 4, 16, 36, 64]

The equivalent loop puts the if inside the for, around the append(): items that fail the test are skipped, so the new list can be shorter than the original.

A conditional expression inside

To choose a value for every item instead of dropping some, put a conditional expression (a if condition else b) in the expression part, at the front. The new list then always has as many items as the original.

temperatures = [14, -2, 21, 0, 8, -5]

labels = ["frost" if t < 0 else "ok" for t in temperatures]
print(labels)

no_negatives = [t if t >= 0 else 0 for t in temperatures]
print(no_negatives)

['ok', 'frost', 'ok', 'ok', 'ok', 'frost']
[14, 0, 21, 0, 8, 0]

Nested comprehensions

A comprehension can contain more than one for clause. They are read from left to right, in the same order as nested loops would be written from top to bottom. The classic use is flattening a list of lists into one list:

weeks = [[12, 14, 11], [15, 17], [9, 10, 13]]

flat = [temp for week in weeks for temp in week]
print(flat)

# the same with loops: the for clauses in the same order
flat = []
for week in weeks:
    for temp in week:
        flat.append(temp)
print(flat)

[12, 14, 11, 15, 17, 9, 10, 13]
[12, 14, 11, 15, 17, 9, 10, 13]

A comprehension can also be the expression of another comprehension. That builds a list of lists, such as a matrix or a grid. Here the inner comprehension builds one row, and the outer one builds a row for every r:

table = [[r * c for c in range(1, 4)] for r in range(1, 4)]
print(table)

matrix = [[1, 2, 3], [4, 5, 6]]
transposed = [[row[i] for row in matrix] for i in range(3)]
print(transposed)

[[1, 2, 3], [2, 4, 6], [3, 6, 9]]
[[1, 4], [2, 5], [3, 6]]

Dictionary and set comprehensions

The same syntax builds other collections: the brackets decide the type of the result.

Dictionary comprehensions

A dictionary comprehension uses curly braces and a key: value pair as its expression. It is the quickest way to build a dictionary from another collection, or to transform or filter an existing one.

names = ["Ada", "Linus", "Grace"]
lengths = {name: len(name) for name in names}
print(lengths)

prices = {"tea": 2.50, "coffee": 3.20, "cake": 4.75}
cheap = {item: price for item, price in prices.items() if price < 4}
print(cheap)

in_cents = {item: round(price * 100) for item, price in prices.items()}
print(in_cents)

{'Ada': 3, 'Linus': 5, 'Grace': 5}
{'tea': 2.5, 'coffee': 3.2}
{'tea': 250, 'coffee': 320, 'cake': 475}

As in any dictionary, keys are unique: if two items produce the same key, the later value replaces the earlier one.

Set comprehensions

A set comprehension uses curly braces with a single expression. The result is a set, so duplicate values are kept only once.

words = ["tea", "cake", "pie", "bread", "jam"]
lengths = {len(word) for word in words}
print(sorted(lengths))

scores = [7, 9, 7, 10, 9, 7]
print(len({s for s in scores}))

[3, 4, 5]
3

Generator expressions

A generator expression looks like a list comprehension in parentheses. It does not build a collection at all: it creates a generator, an object that produces the values lazily, one at a time, only when they are asked for. It therefore uses almost no memory, even for a billion values.

Generator expressions are most useful as the argument of a function that consumes values one by one, such as sum(), min(), max(), any() and all(). When the generator expression is the only argument, its own parentheses may be left out.

print(sum(n * n for n in range(1, 11)))

prices = [2.50, 3.20, 4.75]
print(max(p * 1.1 for p in prices) > 5)
print(any(p > 4 for p in prices))
print(all(p > 4 for p in prices))

385
True
True
False

A generator can be used only once. Each value is produced when requested, and once all values have been produced, the generator is exhausted:

squares = (n * n for n in range(4))

print(next(squares))   # asks for one value
print(list(squares))   # takes all the remaining values
print(list(squares))   # nothing left

0
[1, 4, 9]
[]

Use a list comprehension when you need the values more than once, or need len(), indexing or slicing. Use a generator expression when you go through the values once, especially when there are many.

Readability and limits

Comprehensions are meant to make code easier to read. A short comprehension states "the list of X for every Y" more clearly than four lines of loop. But the benefit disappears quickly when the expression grows. Prefer a plain loop when:

you need more than one for plus a condition, or more than two levels of nesting;
the work per item needs several statements, error handling or temporary variables;
you only want a side effect, such as printing. A list comprehension always builds a list, so [print(x) for x in data] creates a useless list of None values.

orders = [("tea", 2, 2.50), ("cake", 0, 4.75), ("coffee", 3, 3.20)]

# Hard to read in one line:
lines = [f"{name}: {qty * price:.2f}" for name, qty, price in orders if qty > 0]
print(lines)

# Clearer as a loop:
lines = []
for name, qty, price in orders:
    if qty == 0:
        continue
    total = qty * price
    lines.append(f"{name}: {total:.2f}")
print(lines)

['tea: 5.00', 'coffee: 9.60']
['tea: 5.00', 'coffee: 9.60']

The loop variable does not leak

A comprehension has its own scope. Its loop variable exists only inside the brackets and does not overwrite a variable with the same name outside. A for loop is different: its loop variable stays behind with its last value.

n = "unchanged"
squares = [n * n for n in range(5)]
print(n)                  # the comprehension's n was separate

for m in range(5):
    pass
print(m)                  # a for loop's variable remains

unchanged
4