Python List Comprehensions

A comprehension builds a list from a loop in one expression. It is not only shorter, it is also faster than appending in a loop, because the append is done in C.

The parts are always in the same order: the expression, then the for, then any if.

The loop and the comprehension

These two produce the same list.

# the loop and the comprehension do the same thing
squares = []
for n in range(6):
    squares.append(n * n)
print(squares)

print([n * n for n in range(6)])

It prints:

[0, 1, 4, 9, 16, 25]
[0, 1, 4, 9, 16, 25]

Filtering

An if at the end keeps only the items you want.

# filtering with if
print([n for n in range(20) if n % 3 == 0])

It prints:

[0, 3, 6, 9, 12, 15, 18]

Choosing a value

An if/else is part of the expression, so it goes at the front. This one trips people up because the two forms look similar and sit in different places.

# if/else goes before the for
print(['even' if n % 2 == 0 else 'odd' for n in range(5)])

It prints:

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

Rule of thumb: if at the end filters, if/else at the front chooses.

Dict and set comprehensions

Braces give you a dict when you write a pair, and a set when you write a single value.

# dict and set comprehensions
words = ['apple', 'bread', 'apple', 'milk']
print({w: len(w) for w in words})
print({len(w) for w in words})

It prints:

{'apple': 5, 'bread': 5, 'milk': 4}
{4, 5}

The set version dropped the duplicate length, as a set does.

Two for clauses

Multiple for clauses read left to right, in the same order you would nest them in a loop. This is how you flatten a list of lists.

# two for clauses read top to bottom
pairs = [(x, y) for x in 'ab' for y in (1, 2)]
print(pairs)

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

It prints:

[('a', 1), ('a', 2), ('b', 1), ('b', 2)]
[1, 2, 3, 4, 5, 6]

Generator expressions

Round brackets give you a generator instead. It produces values on demand and does not hold the whole list in memory.

# a generator expression does not build the list
import sys

listcomp = [n * n for n in range(100000)]
genexp = (n * n for n in range(100000))
print('list bytes:', sys.getsizeof(listcomp))
print('generator bytes:', sys.getsizeof(genexp))
print('sum is the same:', sum(listcomp) == sum(genexp))

It prints:

list bytes: 800984
generator bytes: 200
sum is the same: True

Same answer, four thousand times less memory. Use a generator when you are passing the result straight to something that consumes it once, such as sum, any or a loop.

The loop variable stays inside

In Python 3 the comprehension has its own scope, so it will not overwrite a variable you already had.

# the loop variable does not leak
n = 'untouched'
result = [n for n in range(3)]
print(result, n)

It prints:

[0, 1, 2] untouched

What to remember

• Order is: expression, for, then if.
• if at the end filters. if/else goes before the for.
• Braces build dicts and sets, round brackets build a generator.
• The loop variable does not leak into the surrounding scope.

If a comprehension needs more than two clauses or will not fit on a line, write the loop. Comprehensions stop paying for themselves as soon as they need to be decoded.