Comprehensions and Generators: Building Sequences Without Building Them

The comprehension forms

print([x * x for x in range(6)])                      # -> [0, 1, 4, 9, 16, 25]
print([x * x for x in range(10) if x % 2 == 0])       # -> [0, 4, 16, 36, 64]

The trailing if filters. To choose between values rather than drop them, use a conditional expression at the front — a different construct in a different position:

print([x if x % 2 else -x for x in range(5)])   # -> [0, 1, -2, 3, -4]

Dict and set comprehensions use braces, distinguished by the presence of a colon:

print({x: x * x for x in range(4)})                  # -> {0: 0, 1: 1, 2: 4, 3: 9}
print(sorted({c for c in "mississippi"}))            # -> ['i', 'm', 'p', 's']
print({v: k for k, v in {"a": 1, "b": 2}.items()})   # -> {1: 'a', 2: 'b'}

Nesting has two distinct shapes and they are easy to confuse. Multiple for clauses read left to right like nested loops and produce a flat result; a comprehension inside a comprehension produces a nested one:

print([(r, c) for r in range(2) for c in range(3)])
# -> [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

print([[r * c for c in range(3)] for r in range(2)])
# -> [[0, 0, 0], [0, 1, 2]]

Flattening is the most useful case of the first shape:

print([y for row in [[1, 2], [3, 4]] for y in row])   # -> [1, 2, 3, 4]

Comprehensions have their own scope, so the loop variable does not leak into the enclosing function — one of the changes Python 3 made over Python 2.

When they hurt. A comprehension earns its place when it maps and filters. It stops earning it when the expression needs a comment, when there are three or more for clauses, when the body has side effects (a comprehension executed only for its side effects is a loop written badly), or when you find yourself using the walrus operator to smuggle in a temporary. The test is whether a reader can state what the result contains after one pass. If not, write the loop.

Generator expressions

Replace the brackets with parentheses and you get a lazy object. Nothing runs until something asks for an item:

g = (x * x for x in range(5))
print(g)                    # -> <generator object <genexpr> at 0x...>
print(next(g), next(g))     # -> 0 1
print(list(g))              # -> [4, 9, 16]     the first two are already consumed

A generator is single-pass: once exhausted it is finished, and iterating it again yields nothing. That is the price of not storing the results.

When a generator expression is the only argument to a function, the parentheses can be omitted:

print(sum(x * x for x in range(5)))              # -> 30
print(any(x > 3 for x in range(5)))              # -> True
print(next((l for l in ["a", "", "b"] if l), None))   # -> a

The last line is a useful idiom: find the first match, or None, without scanning the rest.

The memory difference, measured

One million elements, allocation tracked with tracemalloc on CPython 3.13:

import tracemalloc

tracemalloc.start()
base = tracemalloc.get_traced_memory()[0]
lst = [x * 2 for x in range(1_000_000)]
print(tracemalloc.get_traced_memory()[0] - base)   # -> about 40_400_000

Roughly a factor of \(10^5\), and it is a factor of \(n\) in general: the list is \(O(n)\) in memory, the generator \(O(1)\). Total work is identical if you consume everything — laziness saves memory, not instructions, and it saves time only when you stop early.

So: use a list when you need indexing, len, or more than one pass. Use a generator for a single streaming pass, for data too large to hold, or for a sequence with no end.

yield and generator functions

Any def whose body contains yield becomes a generator function. Calling it does not run the body; it returns a generator. Each next() runs to the following yield, hands back that value, and freezes the frame — locals, instruction pointer and all — until the next request.

def countdown(n):
    while n > 0:
        yield n
        n -= 1
    return "done"

print(list(countdown(3)))   # -> [3, 2, 1]

c = countdown(3)
print(next(c), next(c), next(c))   # -> 3 2 1
next(c)
# StopIteration: 'done'

return inside a generator does not produce a value to the caller; it stops iteration, and its argument is attached to the StopIteration exception (e.value), which is mostly of interest when delegating with yield from.

Because state is preserved rather than accumulated, a generator can be infinite:

import itertools

def fib():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

print(list(itertools.islice(fib(), 10)))
# -> [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

yield from iterable delegates to another iterable, forwarding each of its items — the clean way to compose generators or recurse over a tree:

def gen():
    yield from [1, 2]
    yield 3

print(list(gen()))   # -> [1, 2, 3]

The iterator protocol

for is sugar over two calls. An iterable implements __iter__, which returns an iterator; an iterator implements __next__, returning the next item or raising StopIteration. An iterator must also implement __iter__ returning itself, so that it can be used in a for loop directly.

class Countdown:
    def __init__(self, n):
        self.n = n
    def __iter__(self):
        return self
    def __next__(self):
        if self.n <= 0:
            raise StopIteration
        self.n -= 1
        return self.n + 1

print(list(Countdown(3)))   # -> [3, 2, 1]

The distinction is exactly why a list can be looped over repeatedly and a generator cannot: iter(xs) on a list returns a fresh iterator every time, whereas a generator is its own iterator and hands back itself.

xs = [1, 2, 3]
print(iter(xs) is iter(xs))   # -> False   two independent cursors
i1 = iter(xs)
print(iter(i1) is i1)         # -> True    an iterator returns itself

Generator functions are simply a far shorter way to write that class: the yield statement generates the __next__ machinery, the paused frame is the state, and falling off the end raises StopIteration for you.

itertools highlights

Every one of these returns a lazy iterator, so they compose without materialising intermediates.

results = (x * 2 for x in range(5))
print(sum(results))    # -> 20
print(max(results))    # ValueError: max() iterable argument is empty

That closes the first half of this book. Return to the overview for the map of what comes next.

References

1. Python Software Foundation. List Comprehensions.
2. Python Software Foundation. Functional Programming HOWTO — iterators and generators.
3. Python Software Foundation. itertools — functions creating iterators for efficient looping.
4. Python Software Foundation. Glossary: iterator, iterable, generator.
5. Schemenauer, N., Peters, T., & Hetland, M. L. PEP 255 — Simple Generators.
6. Hettinger, R. PEP 289 — Generator Expressions.
7. van Rossum, G., & Eby, P. J. PEP 380 — Syntax for Delegating to a Subgenerator.