Comprehensions

Write the loop that builds something as one expression — and know when not to.

08of 24
≈ 12Min read
5Parts

The shape

Read it in the order it executes: source, filter, expression. It is not shorthand for its own sake — a comprehension says "I am building a new list from this one", which a for loop with an append only implies.

the same loop, twice
squares = []
for n in numbers:
    if n > 0:
        squares.append(n * n)

squares = [n * n for n in numbers if n > 0]
#          ^^^^^ what to keep
#                ^^^^^^^^^^^^^^ where it comes from
#                               ^^^^^^^^^^ which ones

All four kinds

The generator version is the one worth knowing: it holds no list, produces values as they are asked for, and can run over something that does not fit in memory. When a comprehension is the only argument to a function, the extra parentheses are not needed.

py
[n * n for n in nums]                    # list
{n * n for n in nums}                    # set — duplicates collapse
{name: len(name) for name in names}      # dict
(n * n for n in nums)                    # generator — builds NOTHING yet

sum(n * n for n in nums)                 # the parentheses are implied here
any(w.isupper() for w in words)          # stops at the first True

Nesting reads outside-in

The clauses come in the same order as the nested for statements would. It is the one part of the syntax that people get backwards — write the loop first if you are unsure, then collapse it.

py
flat = [cell for row in grid for cell in row]
# for row in grid:
#     for cell in row:
#         flat.append(cell)

pairs = [(x, y) for x in range(3) for y in range(3) if x != y]

# a nested LIST, rather than a flattened one, is a comprehension inside a comprehension
transposed = [[row[i] for row in grid] for i in range(len(grid[0]))]

The filter goes at the end; the conditional value goes at the front

These two look alike and do different things. if after the for decides whether the item is included; if/else before the for is a conditional expression deciding what the item becomes.

py
[n for n in nums if n > 0]                 # filter: keep some
[n if n > 0 else 0 for n in nums]          # map: change some, keep all
[n for n in nums if n > 0 if n % 2 == 0]   # two filters, and-ed

When not to

Since Python 3 a comprehension has its own scope, so the loop variable does not leak — [x for x in items] leaves no x behind, unlike the for loop it replaces.

  • When it does not fit on a line and a half. Two fors and two ifs is a loop that wants to be a loop.
  • When you are not building anything. A comprehension whose result you throw away is a loop with the wrong punctuation — and it builds a whole list to discard it.
  • When the body needs a statement. try, break and continue cannot go in one; needing them is the signal.
  • When you are folding. sum, min, max, any, all, functools.reduce already exist.

Practice

Next →Functions