Comprehensions

Why it matters

A huge share of everyday code follows one shape: start with an empty collection, loop over some data, maybe skip some items, and add a transformed value for each one. Python has a compact, readable syntax for exactly this, called a comprehensionA one-line syntax that builds a list, dict or set from a loop, optionally filtering items.. You will see comprehensions everywhere in real Python code, especially in data work, so you need to be able to read and write them.

From loop to list comprehension

Here is the familiar loop version, squaring the numbers 1 to 5:

squares = []
for n in range(1, 6):
    squares.append(n * n)
print(squares)    # [1, 4, 9, 16, 25]

The same thing as a list comprehensionA comprehension in square brackets, such as [n * n for n in nums], that builds a list.:

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

Read it left to right as "a list of n * n, for each n in range(1, 6)". The pieces map directly onto the loop: the expression that was inside append(...) comes first, then the for clause.

More examples:

names = ["ada", "grace", "katherine"]
capitalised = [name.title() for name in names]    # ['Ada', 'Grace', 'Katherine']
lengths = [len(name) for name in names]           # [3, 5, 9]

celsius = [0, 21.5, 100]
fahrenheit = [c * 9 / 5 + 32 for c in celsius]    # [32.0, 70.7, 212.0]

texts = ["12.5", "13.0", "11.75"]
values = [float(t) for t in texts]                # [12.5, 13.0, 11.75]

That last one replaces a four-line loop you may have written in the debugging lesson.

Filtering with if

Add an if at the end to keep only some items:

readings = [21.5, -999, 22.0, 23.25, -999, 22.75]

valid = [r for r in readings if r != -999]
print(valid)              # [21.5, 22.0, 23.25, 22.75]

evens = [n for n in range(10) if n % 2 == 0]
print(evens)              # [0, 2, 4, 6, 8]

The loop equivalent puts the if inside the for:

valid = []
for r in readings:
    if r != -999:
        valid.append(r)

If you want to transform values conditionally rather than filter them, put an if ... else expression at the front instead:

labels = ["high" if r > 22 else "normal" for r in valid]
# ['normal', 'normal', 'high', 'high']

Rule of thumb: if at the end filters (fewer items out); if ... else at the front chooses a value (same number of items out).

Dictionary comprehensions

Use braces and a key: value pair:

names = ["ada", "grace", "katherine"]
name_lengths = {name: len(name) for name in names}
# {'ada': 3, 'grace': 5, 'katherine': 9}

scores = {"Ada": 91, "Grace": 48, "Katherine": 85, "Alan": 55}
passed = {name: s for name, s in scores.items() if s >= 50}
# {'Ada': 91, 'Katherine': 85, 'Alan': 55}

prices_usd = {"bolt": 0.10, "nut": 0.05}
prices_cents = {item: round(p * 100) for item, p in prices_usd.items()}
# {'bolt': 10, 'nut': 5}

Set comprehensions

Braces without a colon build a set, which removes duplicates:

emails = ["ada@lab.org", "grace@navy.mil", "alan@lab.org"]
domains = {e.split("@")[1] for e in emails}
# {'lab.org', 'navy.mil'}

Nested loops in a comprehension

You can have more than one for, in the same order as nested loops:

grid = [(row, col) for row in range(2) for col in range(3)]
# [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

matrix = [[1, 2], [3, 4], [5, 6]]
flat = [x for row in matrix for x in row]    # [1, 2, 3, 4, 5, 6]

These get hard to read quickly. If a comprehension needs two for clauses and an if, a regular loop is often kinder to the next reader.

Generator expressions

Swap the square brackets for parentheses inside a function call and you get a generator expressionA comprehension in parentheses that produces values one at a time without building a list, often used inside sum()., which produces values one at a time without building a list in memory:

total = sum(r for r in readings if r != -999)
count_high = sum(1 for r in readings if r > 22)

This is the usual way to sum or count over a filter.

When not to use a comprehension

• When the loop body does more than build one value (printing, updating several variables, handling errors).
• When the line gets long enough that you have to scroll or squint.
• When you do not need the resulting list at all. Do not write [print(x) for x in items]; a plain for loop is clearer.

Common mistakes

• Putting a filtering if at the front. [x if x > 0 for x in xs] is a SyntaxError. Filters go at the end: [x for x in xs if x > 0].
• Forgetting the else in a front if. [a if cond else b for ...] needs both parts.
• Reusing a name. [x for x in x] works but is confusing; use a different loop variable.
• Writing a set by accident. {n for n in nums} is a set, not a dict; a dict needs key: value.

Quick recap

• [expr for item in iterable] builds a list; add if cond at the end to filter.
• {key: value for item in iterable} builds a dict; {expr for ...} builds a set.
• expr_a if cond else expr_b at the front chooses a value for every item.
• sum(... for ...) sums or counts without building a list.
• Prefer a normal loop when a comprehension stops being easy to read.

Try it

In the Do tab you will rewrite three small loops as comprehensions: squares, a filter of passing scores, and a dictionary of word lengths.