Iterators, generators, closures and decorators · about 45 min

Generators

Why this matters

Writing an iterator class by hand, as in the last lesson, takes a dozen lines of boilerplate to remember a position. Generators let you write the same thing as an ordinary-looking function with yield in it, and Python builds the iterator for you. Because generators produce values lazily, one at a time on demand, they let you process files larger than memory, model infinite sequences, and chain processing stages into pipelines where each item flows through all stages before the next item is even read. Much of the standard library (and well-written data-processing code) is built on them.

yield turns a function into a generator

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

for x in countdown(3):
    print(x)
print(list(countdown(5)))
3
2
1
[5, 4, 3, 2, 1]

Compare this with the Countdown class from last lesson: same behavior, a third of the code. Any function containing yield (anywhere in its body) is a generator function. Calling it does not run the body. It returns a generator object, which is an iterator:

g = countdown(2)
print(g)
print(next(g))
print(next(g))
print(next(g))
<generator object countdown at 0x7f...>
2
1
Traceback (most recent call last):
  ...
StopIteration

What happens underneath: a suspended frame

When a normal function is called, Python creates a frame: the object holding its local variables and the position of the instruction being executed. When the function returns, the frame is discarded.

A generator keeps its frame alive between values:

g = countdown(2)     frame created, paused before line 1    (nothing printed)
next(g)              runs until "yield n" with n=2           returns 2, frame paused
next(g)              resumes after the yield: n -= 1,
                     loops, hits "yield n" with n=1          returns 1, frame paused
next(g)              resumes: n -= 1, loop ends,
                     function body finishes                  raises StopIteration

yield is like return that does not destroy the frame: local variables, the loop position, everything stays as it was until the next next(). A return statement (or falling off the end) in a generator raises StopIteration to signal the end. You can see the paused frame in the REPL with g.gi_frame.f_locals and g.gi_frame.f_lineno.

This is why the body runs lazily:

def noisy():
    print("start")
    yield 1
    print("middle")
    yield 2
    print("end")

g = noisy()
print("created")
for v in g:
    print("got", v)
created
start
got 1
middle
got 2
end

Laziness saves memory

A list holds all its elements at once; a generator holds only its current state. Compare:

import sys
squares_list = [n * n for n in range(1_000_000)]
squares_gen = (n * n for n in range(1_000_000))
print(sys.getsizeof(squares_list) > 8_000_000, sys.getsizeof(squares_gen) < 500)
print(sum(squares_gen))
True True
333332833333500000

The list needs about 8 MB just for its pointers (plus the int objects); the generator is a couple of hundred bytes, whatever the range. That is the difference between processing a 20 GB log file line by line and running out of memory.

Generator expressions

The (n * n for n in range(10)) above is a generator expression: comprehension syntax with parentheses, producing a generator instead of a list. When it is the only argument of a function, you can drop the extra parentheses:

words = ["apple", "kiwi", "banana"]
print(sum(len(w) for w in words))
print(max(len(w) for w in words))
print(any(w.startswith("k") for w in words))
15
6
True

Use a list comprehension when you need the list (indexing, several passes, len); use a generator expression when you consume the values once.

Infinite generators

A generator may never finish; the consumer decides when to stop:

import itertools

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

evens = (n for n in naturals() if n % 2 == 0)
print(list(itertools.islice(evens, 5)))
[2, 4, 6, 8, 10]

Each value is computed only when islice asks for it, so the infinite loop costs nothing beyond what is used.

Pipelines

Generators compose. Each stage takes an iterable and yields transformed items:

def read(lines):
    for line in lines:
        line = line.strip()
        if line and not line.startswith("#"):
            yield line

def to_ints(lines):
    for line in lines:
        yield int(line)

def running_total(nums):
    total = 0
    for n in nums:
        total += n
        yield total

data = ["# header", "3", "", "4", "10"]
print(list(running_total(to_ints(read(data)))))
[3, 7, 17]

Data flows through one item at a time: the first line goes through read, to_ints and running_total before the second line is touched. Pass sys.stdin or an open file instead of data and the same code processes any amount of input in constant memory. This is the Unix pipe philosophy inside one program.

yield from

A generator can delegate to another iterable with yield from:

def chain(*iterables):
    for it in iterables:
        yield from it

print(list(chain([1, 2], "ab", range(3))))
[1, 2, 'a', 'b', 0, 1, 2]

yield from it is roughly for x in it: yield x, and it also forwards values sent into the generator and the sub-generator's return value (features used by older coroutine code). It shines in recursive generators such as walking trees or nested lists.

Note: Generators also support send(), throw() and close(), turning them into coroutines that receive values. That machinery was the foundation of asyncio, covered in the concurrency module; for everyday code, yield and yield from are what you need.

Common mistakes

Recap

Going further: explore the itertools module documentation, especially its "Itertools Recipes" section, which is a masterclass in generator pipelines. PEP 255 introduced generators; PEP 380 introduced yield from.

Exercises

Exercise 1 · quiz

When does the body run?

What does this print?

def gen():
    print("A")
    yield 1
    print("B")
    yield 2

g = gen()
print("C")
print(next(g))

Type the printed lines in order, separated by single spaces.

Answer

Calling a generator function runs none of its body: it just creates a generator object (the frame is set up but paused before the first line). The body starts on the first next(), runs until the first yield, hands out 1 and pauses there. B would print only on the second next().

Exercise 2 · write code

Fibonacci generator

In solution.py, write a generator function fib() that yields the Fibonacci numbers forever: 0, 1, 1, 2, 3, 5, 8, ....

Then write first_n(n) that returns a list of the first n Fibonacci numbers, using fib() together with itertools.islice.

first_n(7) returns [0, 1, 1, 2, 3, 5, 8]; first_n(0) returns [].

File: solution.py

Starting point

import itertools


def fib():
    result = []
    a, b = 0, 1
    while True:
        result.append(a)
        a, b = b, a + b
    return result


def first_n(n):
    pass
Hint 1

Replace building a list with yield a: the loop can be infinite because each value is produced on demand.

Hint 2

itertools.islice(fib(), n) stops after n values; wrap it in list(...).

Reference solution
import itertools


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


def first_n(n):
    return list(itertools.islice(fib(), n))

Exercise 3 · write code

A lazy log pipeline

Build a pipeline of generators that processes log lines from standard input. Each line looks like LEVEL message, for example ERROR disk full.

Write three generator functions and connect them:

  1. read_lines(stream): yields each line with surrounding whitespace stripped, skipping blank lines.
  2. parse(lines): yields (level, message) tuples, splitting each line at the first space.
  3. only(level, records): yields only the messages whose level equals level.

Then print each ERROR message, numbered from 1, in the form 1: disk full. After the messages, print errors: N with the count. With no errors, print only errors: 0.

Read from sys.stdin (iterate it directly; do not call .read() or .readlines()).

File: main.py

Starting point

import sys

for line in sys.stdin:
    print(line)

Example

input

INFO started
ERROR disk full

WARN cpu hot
ERROR fan stopped: 0 rpm
INFO done

expected output

1: disk full
2: fan stopped: 0 rpm
errors: 2

Example

input

INFO all good

expected output

errors: 0
Hint 1

Each stage takes an iterable and loops over it, yielding transformed items.

Hint 2

line.partition(" ") returns (before, " ", after) for the first space.

Hint 3

Chain them: only("ERROR", parse(read_lines(sys.stdin))), and use enumerate(..., start=1) for numbering.

Reference solution
import sys


def read_lines(stream):
    for line in stream:
        line = line.strip()
        if line:
            yield line


def parse(lines):
    for line in lines:
        level, _, message = line.partition(" ")
        yield level, message


def only(level, records):
    for lvl, message in records:
        if lvl == level:
            yield message


count = 0
for count, message in enumerate(only("ERROR", parse(read_lines(sys.stdin))), start=1):
    print(f"{count}: {message}")
print(f"errors: {count}")

Exercise 4 · write code

Flatten with yield from

In solution.py, write a recursive generator flatten(items) that yields every non-list value from arbitrarily nested lists, in order, using yield from for the recursive case.

list(flatten([1, [2, [3, 4], []], 5, [[6]]])) is [1, 2, 3, 4, 5, 6]. Only list objects count as nested; strings, tuples and everything else are yielded as single values: list(flatten(["ab", (1, 2)])) is ["ab", (1, 2)].

File: solution.py

Starting point

def flatten(items):
    result = []
    for x in items:
        result.append(x)
    return result
Hint 1

Loop over items; for a list, delegate to a recursive call, otherwise yield x.

Hint 2

yield from flatten(x) yields every value the inner generator produces.

Reference solution
def flatten(items):
    for x in items:
        if isinstance(x, list):
            yield from flatten(x)
        else:
            yield x