8 Comprehensions

A commonly used language construct are comprehensions. They can have performance advantages and can improve readability if properly used.

8.1 List comprehensions

This code of block shows a common pattern that can be simplified:

result = []
for n in range(10):
    if n > 2:
        result.append(n)
result

[3, 4, 5, 6, 7, 8, 9]

To this:

[n for n in range(10) if n>2]

[3, 4, 5, 6, 7, 8, 9]

The logic can also be extended:

result = [n if n>2 else "too-small" for n in range(10)]
result

['too-small', 'too-small', 'too-small', 3, 4, 5, 6, 7, 8, 9]

We can have also nested structures:

results = []
for m in range(2):
    for n in range(3):
        results.append((m,n))
results

[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

results = [(m,n) for m in range(2) for n in range(3)]
results

[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

Tip

Nested comprehension are difficult to read and debug. As a rule of thumb, do not use comprehensions for nested logic beyond 2 levels and, in general, do not add too much logic to comprehensions. If in doubt, just write the loop(s) out.

Although lists are the most commonly used, comprehensions also extend to other collections beyond lists.

8.2 Tuple comprehensions

tuple(n for n in range(10))

(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)

8.3 Set comprehensions

results = set()
for n in range(10):
    if n > 2:
        results.add(n)
results

{3, 4, 5, 6, 7, 8, 9}

results = {n for n in range(10) if n > 2}
results

{3, 4, 5, 6, 7, 8, 9}

8.4 Dictionary comprehensions

results = {}
for n in range(10):
    if n > 2:
        results[f"number_{n}_is"] = n
results

{'number_3_is': 3,
 'number_4_is': 4,
 'number_5_is': 5,
 'number_6_is': 6,
 'number_7_is': 7,
 'number_8_is': 8,
 'number_9_is': 9}

results = {f"number_{n}_is": n for n in range(10) if n > 2}
results

{'number_3_is': 3,
 'number_4_is': 4,
 'number_5_is': 5,
 'number_6_is': 6,
 'number_7_is': 7,
 'number_8_is': 8,
 'number_9_is': 9}

8.5 Exercises

Use this ages dictionary:

ages = {
    "mark": 12,
    "lua": 32,
    "martin": 8,
    "mette": 41,
    "malte": 27,
    "anna": 19,
    "junior": 53,
    "juan": 98,
}

Using a list comprehension, create a new list containing the first letter of each name in ages.
Repeat it, excluding names that start with the letter “a” or “m”.
Repeat it, replacing names starting with “j” with “starts-with-j”, eg “john” -> “starts-with-j”.
Create a set containing the last letter of the each name.
Repeat it, but only consider names with more than 4 letters.
Create a new dictionary using a comprehension that has items (“age_of_NAME”: AGE), reading NAME and AGE from the ages dictionary.