Comprehensions: building a list in one line, and when not to

The three-line loop, written once

squares = []
for n in numbers:
    squares.append(n * n)

squares = [n * n for n in numbers]

Same result. Read it aloud in the order it is written: n times n, for each n in numbers. The expression comes first because the expression is the point; the loop is machinery.

Add a filter on the end:

big = [n for n in numbers if n > 100]
names = [u["name"] for u in users if u["active"]]

And the expression can be anything:

cleaned = [line.strip().lower() for line in lines if line.strip()]

Dictionaries and sets do it too

lengths = {word: len(word) for word in words}
initials = {name[0] for name in names}            # a set

Braces with a colon give a dict comprehension; braces without give a set comprehension. Round brackets give something different, covered below.

A dict comprehension is the standard way to transform a dictionary:

in_rupees = {k: v / 100 for k, v in paise.items()}

Two loops, and the order that trips people

pairs = [(a, b) for a in [1, 2] for b in "xy"]
# [(1, 'x'), (1, 'y'), (2, 'x'), (2, 'y')]

The clauses run left to right, outermost first, exactly as the nested loop would be written. People expect the reverse and get a transposed result.

Flattening a nested list is the common use:

flat = [item for row in grid for item in row]

Two levels is the honest limit. Three, or two plus conditions, and a plain loop is easier for the next person — including you in a month.

Where the if goes changes the meaning

[n if n > 0 else 0 for n in numbers]     # replaces negatives with 0 — same length
[n for n in numbers if n > 0]            # drops negatives — shorter

A conditional before the for is a choice of value and keeps every element. A plain if after the for is a filter and removes elements. Both are valid, they look almost identical, and they do different things. When a comprehension returns the wrong number of items, this is why.

Round brackets give a generator, not a tuple

total = sum(n * n for n in numbers)

That is a generator expression. It produces values one at a time rather than building a list, so sum never holds a million squares in memory at once. For a large input the difference is the difference between 8 MB and 80 bytes.

There is no tuple comprehension. tuple(n for n in xs) is how you get one.

Generators are single-use: once consumed, they are empty. Assigning one to a variable and looping over it twice gives you results the first time and nothing the second, with no error. The generators lesson later in the course takes this apart properly.

Is it faster?

Slightly. A list comprehension avoids a method lookup and a function call per item, so it typically runs 20–30 per cent faster than the equivalent append loop. That is worth having and it is not a reason to use one. The reason to use one is that it says "this list is built from that list", in one line, with no chance of appending to the wrong variable.

When not to use one

• When it does something rather than produces something. [print(x) for x in items] builds a list of None values and throws it away. Write the loop.
• When it needs a try. Comprehensions cannot catch exceptions. If a conversion may fail, loop.
• When it will not fit on a line and a half. A comprehension spanning four lines with two conditions is harder to read than the loop it replaced, and reviewers will say so.
• When you need the intermediate values for debugging. A loop lets you print inside it; a comprehension does not.

The scope detail

In Python 3, the loop variable of a comprehension is local to it:

n = "untouched"
squares = [n * n for n in range(3)]
print(n)      # 'untouched'

In Python 2 that variable leaked out and overwrote n. Old code sometimes depends on the leak, which is one more reason not to run Python 2.

Try this now

Take a list of dictionaries with name and marks. Produce, each in one line: the names of everyone above 60, a dictionary of name to marks, and the average marks using a generator expression inside sum.