Python: generators

A function that pauses itself#

An ordinary function runs until return and forgets its state. A function with yield behaves differently: it hands out a value and freezes, keeping everything — where it stopped and its local variables — until the next request.

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

gen = countdown(3)
print(gen)          # <generator object countdown at 0x...>
print(next(gen))    # 3
print(next(gen))    # 2
print(next(gen))    # 1
next(gen)           # prints "done", then StopIteration

Calling countdown(3) does not run a single line of the body — it only creates a generator object. The first line runs on the first next. When the body reaches its end, the generator raises StopIteration; a for loop catches that exception itself, so you do not see it in ordinary code:

for x in countdown(3):
    print(x, end=" ")   # 3 2 1 done

Generator expressions#

If the logic fits in one expression, you do not need a separate function. The syntax is the same as a list comprehension, but in parentheses:

squares_list = [x * x for x in range(5)]      # a list, computed in full
squares_gen = (x * x for x in range(5))       # a generator, nothing computed yet
print(squares_list)                            # [0, 1, 4, 9, 16]
print(squares_gen)                             # <generator object ...>
print(sum(x * x for x in range(5)))            # 30

When a generator expression is the only argument of a function, you can drop the extra parentheses: sum(x * x for x in range(5)). This is the most common use: computing a sum, a maximum or any/all without building a list.

lines = ["1", "", "3"]
print(any(line == "" for line in lines))    # True
print(max(len(line) for line in lines))     # 1

Why you need them#

Memory. A list comprehension over a million items creates a million objects at once; a generator holds one. You can see the difference with sys.getsizeof: a list's size grows with its length, a generator's stays constant.

Early exit. If the answer is found at the third item, the rest are never computed:

def logged(items):
    for item in items:
        print("computing", item)
        yield item * 2

for value in logged([1, 2, 3]):
    if value >= 4:
        break
print("stopped")
  # computing 1 / computing 2 / stopped

Infinity. A generator may never end — that is fine as long as something else limits it:

def naturals():
    n = 1
    while True:
        yield n
        n += 1

from itertools import islice
print(list(islice(naturals(), 5)))   # [1, 2, 3, 4, 5]

yield from passes along every value of another generator or sequence without an explicit loop: yield from range(3) yields 0, 1, 2.

Pitfalls#

A generator is single-use. Once you have gone through it, it is empty:

gen = (x for x in [1, 2, 3])
print(list(gen))   # [1, 2, 3]
print(list(gen))   # []

This bites in code that first prints the result and then tries to count it. If you need to process the data again, save it to a list or create the generator anew.

len does not work. A generator has no length: len(gen) raises TypeError: object of type 'generator' has no len(). You can count items only by going through them: sum(1 for _ in gen) — and after that the generator is empty.

No indexing either. gen[0] is impossible; use next(gen) or itertools.islice.

Late binding. A generator expression reads outer variables at the moment it is iterated, not when it is created. If a variable changes in between, the result changes too.

A short exercise: write a generator read_numbers(path) that yields one number at a time from a file, and add them up with sum(...) without building a list. Reading line by line is covered in the article on reading files in Python, and how functions with parameters work is in the article on Python functions.

Step-by-step plan

1. Your first yieldWrite a generator of three values and call next four times until you get StopIteration.
2. Compare with a listBuild the same data with a list comprehension and a generator expression and print both objects.
3. Count without a listFind the sum of squares of the first thousand numbers using sum with a generator expression.
4. See the lazinessPut a print inside the generator and break out of the loop early — check that nothing extra was computed.
5. Hit the single-use trapLoop over a generator twice in a row and explain why the second pass is empty.

Check yourself

Sources

• Generators in the Python tutorialThe Iterators, Generators and Generator Expressions sections
  free
• The itertools moduleislice, count, chain and other ready-made generators
  free
• Python glossaryExact definitions of iterator, generator and iterable
  free

Was this helpful?

Yes No