Comprehensions

What you'll learn

• Build a list from another collection in one expression with a list comprehension
• Filter with an if clause and transform with any expression
• Write dictionary and set comprehensions, and know what a generator expression is
• Decide when a comprehension is clearer than a loop, and when it is not

Concept

You have written this loop many times: create an empty list, loop over something, append a transformed item. A list comprehension is that loop folded into one expression: [price * 1.2 for price in prices]. It reads left to right as "a list of price times 1.2, for each price in prices", and it produces exactly what the four-line loop did. Python programmers use comprehensions constantly, so reading them fluently matters as much as writing them.

A comprehension has three parts. The expression at the front decides what goes into the result. The for clause names the source. An optional if clause at the end filters, keeping only items where it is true. Read them in that order when writing, and in the order for, if, expression when reading: first what you loop over, then what you keep, then what you make of it.

The same shape builds dictionaries with {key: value for ...} and sets with {item for ...}. Parentheses instead of brackets give a generator expression, which produces items one at a time instead of building a list; you have already used those inside sum(...), any(...), and max(...). Generators get their own lesson later. For now: brackets when you need the list, parentheses when you only need to loop once.

Comprehensions are for building a collection from another collection. They are not for running code with side effects, and they stop being clear the moment they need a second for or a nested if/else that does not fit on one line. The rule: if you cannot read it aloud as one sentence, write the loop. A clear four-line loop beats a clever one-liner every time someone else has to maintain it.

Code

The loop and the comprehension side by side. They produce identical lists; the second is the one you will write from now on.

# same_result.py
prices = [3.5, 12.0, 7.25, 0.99]

with_tax_loop = []
for price in prices:
    with_tax_loop.append(round(price * 1.2, 2))

with_tax = [round(price * 1.2, 2) for price in prices]

print(with_tax_loop)
print(with_tax)
print(with_tax_loop == with_tax)

[4.2, 14.4, 8.7, 1.19]
[4.2, 14.4, 8.7, 1.19]
True

Filtering and transforming together, plus a conditional expression inside the expression part. Note the difference: the trailing if drops items, the if/else at the front changes them.

# filter_transform.py
scores = [72, 95, 58, 88, 61, 45]

passed = [s for s in scores if s >= 60]
labels = ["pass" if s >= 60 else "fail" for s in scores]
curved = [min(s + 5, 100) for s in scores if s < 90]

print(passed)
print(labels)
print(curved)

[72, 95, 88, 61]
['pass', 'pass', 'fail', 'pass', 'pass', 'fail']
[77, 63, 93, 66, 50]

Dictionary and set comprehensions, and a generator expression feeding sum. The set removes duplicate first letters automatically.

# dict_set_gen.py
words = ["apple", "avocado", "banana", "blueberry", "cherry"]

lengths = {w: len(w) for w in words}
first_letters = {w[0] for w in words}
total_letters = sum(len(w) for w in words)

print(lengths)
print(sorted(first_letters))
print(total_letters)

{'apple': 5, 'avocado': 7, 'banana': 6, 'blueberry': 9, 'cherry': 6}
['a', 'b', 'c']
33

Where the line is. The first comprehension flattens a grid and is still readable; the second does the same job with a nested condition and is not, so the loop underneath it is the better version.

# readable_limit.py
grid = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

flat = [n for row in grid for n in row]
print(flat)

hard = [n * 2 if n % 2 == 0 else -n for row in grid for n in row if n != 5]
clear = []
for row in grid:
    for n in row:
        if n == 5:
            continue
        clear.append(n * 2 if n % 2 == 0 else -n)
print(hard == clear, clear)

[1, 2, 3, 4, 5, 6, 7, 8, 9]
True [-1, 4, -3, 8, 12, -7, 16, -9]

Try it yourself

Create comp.py with a list of ten temperatures in Celsius. Build a list of them in Fahrenheit, a list of only the days above 25 degrees, a dictionary from day number to temperature, and a set of the temperatures rounded to the nearest 5. Then take one of the comprehensions and rewrite it as a loop to confirm they match.

Common mistakes

• Using a comprehension for side effects. [print(x) for x in items] builds a useless list of Nones; a for loop is the tool for doing things.
• Confusing the two ifs. [x for x in xs if cond] filters; [a if cond else b for x in xs] transforms; the first has no else.
• Nesting until it is unreadable. Two for clauses is the practical limit; beyond that, write the loop.

What's next

• A comprehension is a loop that builds a collection, folded into one expression: expression, for source, optional if filter.
• Brackets make a list, braces make a dictionary or set, parentheses make a generator for one-pass use.
• Use them for one readable transformation; when the line needs explaining, write the loop.

You can build collections in one line. Next you will move data in and out of your programs, with files and paths.