List Comprehensions: Syntax and Use

A list comprehension builds a whole new list in a single expression. The full form is [expression for item in iterable if condition], and you read it from left to right as "take each item, test it, and work out a value for it".

numbers = [1, 2, 3, 4, 5, 6]

squares = [n * n for n in numbers]
print(squares)          # [1, 4, 9, 16, 25, 36]

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

labels = ['even' if n % 2 == 0 else 'odd' for n in numbers]
print(labels)           # ['odd', 'even', 'odd', 'even', 'odd', 'even']

Three different things are happening in that example. The first comprehension maps every item through an expression, so each number is squared and the result is a list of six values. The second one adds an if clause, which filters. The items that fail the test never appear in the output, which is why the result has four elements instead of six. The third one uses a conditional expression inside the comprehension, and that chooses a value for every single item without dropping any of them.

The difference between the if clause and the conditional expression is worth being able to explain, because the two look similar and they do completely different jobs. The if clause decides whether an item appears in the result at all. The conditional expression decides what value an item that is already included will hold. A comprehension is allowed to use both of them together.

There are two more details that trip people up. The expression comes first and the for clause comes after it, which is the opposite of the set notation that most people meet at school. And because the whole thing is one expression, a comprehension can be assigned to a variable, returned from a function, passed as an argument, or placed inside another comprehension.

The if clause filters, and more than one for clause gives you nested loops inside a single expression.

# Nested loops. The outer one runs first.
pairs = [(a, b) for a in [1, 2] for b in ['x', 'y']]
print(pairs)          # [(1, 'x'), (1, 'y'), (2, 'x'), (2, 'y')]

# Flattening a grid.
matrix = [[1, 2], [3, 4]]
flat = [value for row in matrix for value in row]
print(flat)           # [1, 2, 3, 4]

# More than one condition means both of them must pass.
print([n for n in range(10) if n % 3 == 0 if n > 2])   # [3, 6, 9]

The order of the for clauses runs from left to right, and that order decides the order of the output. In the pairs comprehension the outer loop walks the numbers, so every value of b is produced while a is 1, and only then does the comprehension move on to a being 2.

The flattening comprehension is the pattern that is worth remembering. The outer for walks the rows and the inner for walks the items inside each row, so the values come out in row order. Swapping the two clauses would not work, because a row is not a number and you cannot do that kind of thing with it.

More than one if clause is allowed, and it means that all of them must be satisfied. The last line keeps the multiples of 3 that are also greater than 2, which gives 3, 6, and 9.

Two practical notes are worth remembering. The condition is checked once for every item, so the order of the clauses does not save you any work. And when a comprehension grows past two or three lines, it is usually clearer to write it as an ordinary loop.

The syntax is [expression for item in iterable if condition]. The expression is worked out for every item that passes the condition, and all of those results are collected into a new list. Because the whole thing is an expression, it can be assigned, returned, or nested inside another expression.

Two rules govern the order. The expression comes before the for clause, which is the reverse of the set notation that most people meet at school. And when there is more than one for clause, the leftmost one is the outermost loop, which is exactly the same left-to-right nesting that you get with nested for statements.

There are four good reasons to avoid a comprehension. Use a normal loop when the body needs statements instead of expressions, because a comprehension cannot hold an assignment, a try block, or a break. Use a loop when the comprehension would have to wrap over several lines, because it becomes harder to read than a loop. Use a loop when the expression needs a comment to explain it, because a comprehension cannot carry a comment. And use a generator expression instead when you are building a very large list, because the generator does not hold everything in memory at once.

That last point is the comparison with generator expressions. Square brackets build the whole list immediately, while round brackets build a generator that produces its items one at a time, as they are needed. A comprehension is the right choice when you genuinely need the list, such as when you will index into it, walk through it more than once, or ask for its length. A generator expression is the right choice when the results are consumed once in a pipeline, where the intermediate list would be pure waste.

One more thing is worth saying in an interview. A comprehension is usually faster than the equivalent loop, because it compiles to a dedicated bytecode loop instead of a Python-level for statement. That difference is only a constant factor though, so readability should win when the two versions are close.