Comprehensions ✨

From loop to comprehension

Here is the loop you already know:

squares = []
for n in range(1, 6):
    squares.append(n * n)
print(squares)

[1, 4, 9, 16, 25]

And here it is again:

squares = [n * n for n in range(1, 6)]
print(squares)

[1, 4, 9, 16, 25]

The pieces are the same, rearranged. The thing you append comes first, then the loop. Read it aloud: "n times n, for each n in one to five".

[  n * n     for n in range(1, 6)  ]
   ^^^^^     ^^^^^^^^^^^^^^^^^^^^
   what      where it comes from
   you keep

Adding a filter

numbers = [4, 9, 15, 22, 30, 7]

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

big_doubled = [n * 2 for n in numbers if n > 10]
print(big_doubled)

[4, 22, 30]
[30, 44, 60]

The if goes at the end and decides what gets in. Reading order:

1. Take each n from numbers,
2. keep it only if the condition is true,
3. and put n * 2 in the new list.

Transforming text

crew = ["  guybrush ", "ELAINE", " otis  "]

tidy = [name.strip().title() for name in crew]
print(tidy)

lengths = [len(name) for name in tidy]
print(lengths)

initials = [name[0] for name in tidy]
print(initials)

['Guybrush', 'Elaine', 'Otis']
[8, 6, 4]
['G', 'E', 'O']

Dictionary and set comprehensions

words = ["grog", "sword", "map"]

lengths = {word: len(word) for word in words}
print(lengths)

first_letters = {word[0] for word in words}
print(sorted(first_letters))

scores = {"Guybrush": 95, "Elaine": 88, "Otis": 42}
passed = {name: score for name, score in scores.items() if score >= 50}
print(passed)

{'grog': 4, 'sword': 5, 'map': 3}
['g', 'm', 's']
{'Guybrush': 95, 'Elaine': 88}

Same syntax, different brackets. [ ] makes a list, { } with a colon makes a dictionary, { } without makes a set. Filtering a dictionary down to the entries you want is one of the most useful lines in day-to-day Python.

The conditional version

numbers = [4, 9, 15, 22]

labels = ["even" if n % 2 == 0 else "odd" for n in numbers]
print(labels)

['even', 'odd', 'odd', 'even']

Nested loops in a comprehension

pairs = [(a, b) for a in "AB" for b in [1, 2]]
print(pairs)

grid = [[1, 2], [3, 4], [5, 6]]
flat = [n for row in grid for n in row]
print(flat)

[('A', 1), ('A', 2), ('B', 1), ('B', 2)]
[1, 2, 3, 4, 5, 6]

The loops read in the same order you would write them normally: outer first, inner second. Flattening a list of lists is the common case and worth memorising as a phrase.

Generator expressions: the lazy cousin

numbers = range(1, 1_000_001)

# builds a million-item list in memory first
total_list = sum([n * n for n in numbers])

# produces one value at a time and never stores them all
total_gen = sum(n * n for n in numbers)

print(total_list == total_gen)
print(f"{total_gen:,}")

True
333,333,833,333,500,000

Round brackets (or none at all, inside a function call) give you a generator expression: it computes values on demand instead of building the whole list. For a million squares that is the difference between using about 40MB of memory and using almost none. When you are feeding straight into sum, max, any or a for loop, drop the square brackets. Lesson 34 goes deeper.

any and all

scores = [95, 88, 42, 71]

print(any(s < 50 for s in scores))
print(all(s >= 40 for s in scores))
print(sum(1 for s in scores if s >= 70))

True
True
3

any and all plus a generator expression is how you ask "is there at least one" and "is every single one" in a single readable line. They also stop early as soon as the answer is decided.

When not to use one

# Do not do this to people.
result = [y for x in range(10) if x % 2 == 0 for y in range(x) if y % 3 == 0 and y > 1]
print(result)

[3, 3, 3, 6]

Also: if the loop is doing something rather than collecting something, use a loop.

# wrong: building a list of Nones purely for the side effect
# [print(n) for n in range(3)]

# right
for n in range(3):
    print(n)

0
1
2

# A
result = []
for word in ["grog", "map", "sword"]:
    result.append(word.upper())

# B
long_words = []
for word in ["grog", "a", "sword", "of"]:
    if len(word) > 2:
        long_words.append(word)

# C
lengths = {}
for word in ["grog", "map"]:
    lengths[word] = len(word)

result = [word.upper() for word in ["grog", "map", "sword"]]
long_words = [w for w in ["grog", "a", "sword", "of"] if len(w) > 2]
lengths = {word: len(word) for word in ["grog", "map"]}

print(result)
print(long_words)
print(lengths)

['GROG', 'MAP', 'SWORD']
['grog', 'sword']
{'grog': 4, 'map': 3}

result = [
    "FizzBuzz" if n % 15 == 0
    else "Fizz" if n % 3 == 0
    else "Buzz" if n % 5 == 0
    else str(n)
    for n in range(1, 16)
]
print(result)

['1', '2', 'Fizz', '4', 'Buzz', 'Fizz', '7', '8', 'Fizz', 'Buzz', '11', 'Fizz', '13', '14', 'FizzBuzz']

crew = [
    {"name": "Guybrush", "pay": 100},
    {"name": "Otis", "pay": 40},
    {"name": "Elaine", "pay": 250},
    {"name": "Meathook", "pay": 45},
]

crew = [
    {"name": "Guybrush", "pay": 100},
    {"name": "Otis", "pay": 40},
    {"name": "Elaine", "pay": 250},
    {"name": "Meathook", "pay": 45},
]

well_paid = sorted(m["name"] for m in crew if m["pay"] > 50)
print(well_paid)

['Elaine', 'Guybrush']