List and Set Comprehensions

The shape

[ expression   for item in iterable   if condition ]
   ^              ^                     ^
   what to keep   where it comes from   which ones (optional)

# The loop
squares = []
for n in range(6):
    squares.append(n * n)

# The comprehension
squares = [n * n for n in range(6)]

print(squares)      # [0, 1, 4, 9, 16, 25]

Read it aloud as "n times n, for each n in range 6". The for part is the same loop header you would have written; the expression at the front is what would have gone inside append.

Filtering

numbers = [4, -2, 7, 0, -5, 9]

positives = [n for n in numbers if n > 0]
print(positives)                     # [4, 7, 9]

evens = [n for n in range(20) if n % 2 == 0]
print(evens)

words = ["apple", "fig", "banana", "kiwi"]
long_words = [w for w in words if len(w) > 4]
print(long_words)                    # ['apple', 'banana']

# Two conditions
print([n for n in range(30) if n % 3 == 0 if n % 5 == 0])    # [0, 15]
print([n for n in range(30) if n % 3 == 0 and n % 5 == 0])   # the same thing

Transforming

names = ["  meera ", "ARUN", "Sara  "]

cleaned = [name.strip().title() for name in names]
print(cleaned)                       # ['Meera', 'Arun', 'Sara']

lengths = [len(name) for name in cleaned]
print(lengths)                       # [5, 4, 4]

pairs = [(name, len(name)) for name in cleaned]
print(pairs)                         # [('Meera', 5), ('Arun', 4), ('Sara', 4)]

celsius = [0, 20, 37, 100]
fahrenheit = [c * 9 / 5 + 32 for c in celsius]
print(fahrenheit)                    # [32.0, 68.0, 98.6, 212.0]

Filter and transform together

raw = ["12", "abc", "7", "", "34"]

numbers = [int(value) for value in raw if value.isdigit()]
print(numbers)                       # [12, 7, 34]

records = [
    {"name": "Meera", "score": 92},
    {"name": "Arun", "score": 45},
    {"name": "Sara", "score": 78},
]

passed = [r["name"] for r in records if r["score"] >= 50]
print(passed)                        # ['Meera', 'Sara']

Conditional expressions inside a comprehension

There are two different places a condition can appear, and they do different jobs.

numbers = [1, -2, 3, -4]

# if at the END filters: some items are dropped
print([n for n in numbers if n > 0])                  # [1, 3]

# if...else at the FRONT transforms: every item is kept
print([n if n > 0 else 0 for n in numbers])           # [1, 0, 3, 0]

# Both together
print([n * 2 if n > 0 else 0 for n in numbers if n != -4])   # [2, 0, 6]

Note that the front position requires an else, because it is a conditional expression and every expression must produce a value.

Nested loops

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

# The loops read in exactly the same order as nested for statements:
pairs = []
for x in [1, 2]:
    for y in ["a", "b"]:
        pairs.append((x, y))

Flattening

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

flat = [value for row in matrix for value in row]
print(flat)       # [1, 2, 3, 4, 5, 6]

Outer loop first, inner loop second. Writing them the other way round is the most common comprehension error, and it produces a NameError because row is not yet defined.

Nested comprehensions

grid = [[0] * 3 for _ in range(2)]         # a fresh row each time
print(grid)                                 # [[0, 0, 0], [0, 0, 0]]

matrix = [[1, 2], [3, 4]]
doubled = [[value * 2 for value in row] for row in matrix]
print(doubled)                              # [[2, 4], [6, 8]]

transposed = [[row[c] for row in matrix] for c in range(len(matrix[0]))]
print(transposed)                           # [[1, 3], [2, 4]]

A comprehension whose expression is itself a comprehension builds a nested structure. This is the correct way to build a grid, because each inner comprehension creates a new list.

Set comprehensions

words = ["apple", "avocado", "banana", "blueberry"]

initials = {w[0] for w in words}
print(initials)                    # {'a', 'b'}

print({n % 5 for n in range(20)})  # {0, 1, 2, 3, 4}
print({len(w) for w in words})     # the distinct lengths

sentence = "the cat sat on the mat"
print(sorted({w for w in sentence.split() if len(w) == 3}))

The only change is the brackets. Duplicates collapse automatically, and the result has no order.

Generator expressions

squares_list = [n * n for n in range(1_000_000)]      # builds a million items
squares_gen = (n * n for n in range(1_000_000))       # builds nothing yet

print(sum(n * n for n in range(1_000_000)))           # no list is ever created

Round brackets produce a generator, which computes values one at a time as they are needed. When the result is fed straight into sum, max, any, all, join or a for loop, use a generator and save the memory.

names = ["Meera", "Arun"]

print(", ".join(n.upper() for n in names))       # no brackets needed inside a call
print(any(len(n) > 4 for n in names))            # True
print(max((len(n) for n in names), default=0))   # 5

Comprehension or loop?

# Good: one clear transformation
names = [r["name"] for r in records if r["active"]]

# Bad: too much happening
result = [transform(x) for sublist in data for x in sublist
          if check(x) and x.value > threshold and x.name not in excluded]

A comprehension is right when it builds one collection from one source with at most one filter. When you need several statements, error handling, logging or a running total, use a loop. Readability is the deciding factor, not cleverness.

# A comprehension cannot do this - use a loop
results = []
for value in raw_values:
    try:
        results.append(int(value))
    except ValueError:
        print("skipping", value)

Scope

n = "outer"
squares = [n * n for n in range(3)]
print(n)          # outer - the comprehension variable does NOT leak

# Compare with a plain loop
for n in range(3):
    pass
print(n)          # 2 - the loop variable DOES survive

Common mistakes

• Writing the nested for clauses in the wrong order.
• Forgetting the else in a front position conditional expression.
• Building a large list where a generator would do.
• Using a comprehension purely for a side effect, such as calling print inside it. Use a loop.
• Using [[0] * n] * m instead of a comprehension for a grid.
• Packing three conditions and two loops into one line.

Best practices

• Keep a comprehension to one line, or split it across lines at the for and if.
• Use a generator expression whenever the result is consumed immediately.
• Use a loop for anything involving try, several statements, or accumulation.
• Name the result meaningfully; the comprehension explains the how, the name explains the what.

Practice

1. Build a list of the squares of the odd numbers below 30.
2. Given a list of sentences, produce the set of all distinct words longer than three characters.
3. Replace every negative number in a list with zero, keeping the list length the same.
4. Flatten a list of lists and explain why reversing the two for clauses fails.
5. Rewrite a comprehension that must skip invalid values with error handling, as a loop, and explain why.

Conclusion

A comprehension is a loop turned inside out: the result first, then where it comes from, then which items qualify. Use it for one clear transformation, use a generator when the result is consumed straight away, and drop back to a loop the moment it stops reading like a sentence.