Comprehensions¶

~20 min · prerequisite: Lists

What you will be able to do¶

Build a new list from an existing collection with a list comprehension, optionally keeping only some items with if.
Build a dictionary or a set in one expression with a dictionary or set comprehension.
Recognize when a plain for loop is clearer than a comprehension.

The idea¶

Lesson 2 built lists with a loop and append. A comprehension writes that pattern in one expression inside brackets.

The general form of a list comprehension is [expression for name in collection]. For each item in collection, Python stores the item in name, evaluates expression, and adds the result to a new list. The loop [n * n for n in range(1, 6)] gives [1, 4, 9, 16, 25]. Read the brackets from left to right: first the expression, then the loop that supplies its values. The original collection is never changed, because a comprehension always creates a new one.

A comprehension can filter. Add if condition at the end, and only items for which the condition is True produce a result. The form [n for n in nums if n % 2 == 0] keeps the even numbers. The % operator gives the remainder, as you saw in the arithmetic lesson of Part 1.

A dictionary comprehension uses curly braces with a colon: {key_expr: value_expr for name in collection}. A set comprehension uses curly braces with one expression, and duplicates disappear because a set holds each value once. The kind of result depends only on the brackets and the form inside them, so you choose the result type by choosing the brackets.

The loop variable belongs to the comprehension. The name n does not exist after the line finishes, so you cannot read it afterward. The new collection is the only result.

Use a comprehension for a simple transformation. When the expression needs several steps, a for loop is easier to read.

Worked example¶

numbers = [1, 2, 3, 4, 5, 6]

squares = [n * n for n in numbers]
print(squares)

evens = [n for n in numbers if n % 2 == 0]
print(evens)

words = ["apple", "fig", "banana"]
lengths = {word: len(word) for word in words}
print(lengths)

unique_lengths = {len(word) for word in words}
print(sorted(unique_lengths))

Line by line, this is what Python does:

[n * n for n in numbers] squares each number in turn and collects the squares in a new list.
[n for n in numbers if n % 2 == 0] keeps only the even numbers, which gives [2, 4, 6].
{word: len(word) for word in words} builds a dictionary. Each word becomes a key, and its length becomes the value.
{len(word) for word in words} builds a set of the lengths 5, 3, and 6. sorted puts them in order for printing.

The output is:

[1, 4, 9, 16, 25, 36]
[2, 4, 6]
{'apple': 5, 'fig': 3, 'banana': 6}
[3, 5, 6]

Try it¶

The first block checks a list comprehension with and without a filter. The second block checks dictionary and set comprehensions, and the nested two-loop form.

nums = [1, 2, 3, 4, 5]
assert [n * 10 for n in nums] == [10, 20, 30, 40, 50]
assert [n for n in nums if n > 3] == [4, 5]
assert [n for n in nums if n % 2 == 1] == [1, 3, 5]
assert nums == [1, 2, 3, 4, 5]

words = ["cat", "dog", "cat", "emu"]
lengths = {w: len(w) for w in words}
assert lengths == {"cat": 3, "dog": 3, "emu": 3}
letters = {w[0] for w in words}
assert letters == {"c", "d", "e"}
assert len(letters) == 3
pairs = [(x, y) for x in [1, 2] for y in ["a", "b"]]
assert pairs == [(1, "a"), (1, "b"), (2, "a"), (2, "b")]

Common mistakes¶

Reading the loop variable after the comprehension. The code squares = [n * n for n in [1, 2, 3]] followed by print(n) gives:
```
NameError: name 'n' is not defined
```
The variable n exists only inside the comprehension. Store the result in squares and use that.
Putting the filter before the loop. The code [n if n > 2 for n in [1, 2, 3]] gives:
```
SyntaxError: invalid syntax
```
Python reads an if placed before for as a choice between two values, so it needs an else. Put the condition after the loop, as in [n for n in [1, 2, 3] if n > 2], to keep only the items that pass the test.
Putting a list in a set comprehension. The code {[w] for w in ["a"]} gives:
```
TypeError: unhashable type: 'list'
```
Set items must be hashable. Use the string itself, or a tuple such as (w,).

Practice¶

Create nums = [2, 4, 6]. Use a list comprehension to build halves with each number divided by 2 using /. Assert that halves equals [1.0, 2.0, 3.0].
Create words = ["hi", "hello", "hey"]. Use a list comprehension to keep only the words with length 3 or more. Assert that the result equals ["hello", "hey"].
Create temps = [12, 25, 18]. Use a dictionary comprehension to map each temperature to "warm" if it is 20 or more, and to "cool" otherwise. Assert the result equals {12: "cool", 25: "warm", 18: "cool"}.
Create letters = ["a", "b", "a", "c"]. Use a set comprehension to build upper_set of the uppercase letters. Assert it equals {"A", "B", "C"}.
Stretch: build grid, a list of two rows, where row r holds r * 3 + c for each c in range(3). Use a comprehension inside another comprehension. Assert that grid equals [[0, 1, 2], [3, 4, 5]].

Solution

python nums = [2, 4, 6] halves = [n / 2 for n in nums] assert halves == [1.0, 2.0, 3.0]
python words = ["hi", "hello", "hey"] long_words = [w for w in words if len(w) >= 3] assert long_words == ["hello", "hey"]
python temps = [12, 25, 18] labels = {t: ("warm" if t >= 20 else "cool") for t in temps} assert labels == {12: "cool", 25: "warm", 18: "cool"}
python letters = ["a", "b", "a", "c"] upper_set = {ch.upper() for ch in letters} assert upper_set == {"A", "B", "C"}
python grid = [[row * 3 + col for col in range(3)] for row in range(2)] assert grid == [[0, 1, 2], [3, 4, 5]]

Recap¶

A list comprehension, [expr for name in collection if condition], builds a new list in one expression.
Dictionary and set comprehensions use curly braces. A dictionary comprehension has key: value, and a set comprehension removes duplicates.
The loop variable exists only inside the comprehension.