Python · Comprehensions

A comprehension builds a collection in one expression. It reads in the same order as the loop it replaces:

[expression for item in iterable if condition]

• {k: v for ...} builds a dict, and {x for ...} builds a set.
• You can chain loops: [(i, j) for i in a for j in b]. The loops nest left to right.
• (x for ...) in round brackets is a generator expression. It’s lazy, so sum(x*x for x in range(10**6)) never builds the full list.

My rule from today: if the comprehension needs more than one if, or a nested loop that is hard to read, write the plain loop instead.

squares = [n * n for n in range(10)]
evens = [n for n in squares if n % 2 == 0]
print(squares)
print(evens)

words = ["python", "numpy", "pandas", "torch"]
lengths = {w: len(w) for w in words}
firsts = {w[0] for w in words}
print(lengths)
print(sorted(firsts))

pairs = [(i, j) for i in range(3) for j in range(3) if i < j]
print(pairs)

matrix = [[1, 2, 3], [4, 5, 6]]
transpose = [[row[c] for row in matrix] for c in range(3)]
print(transpose)
print(sum(n * n for n in range(1_000_000)))

Output:

[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
[0, 4, 16, 36, 64]
{'python': 6, 'numpy': 5, 'pandas': 6, 'torch': 5}
['n', 'p', 't']
[(0, 1), (0, 2), (1, 2)]
[[1, 4], [2, 5], [3, 6]]
333332833333500000