Generators

Concepts covered: Generators

Work that waits to be asked

f.read() returns the whole file as 1 string before any work starts, so the memory bill equals the file's size. for line in f pays for 1 line at a time. A generator is how you write a function with that second behavior: it produces values 1 at a time when asked and holds only what it needs to produce the next one.

yield pauses a function

A function that contains yield is a generator function. Calling it runs none of its lines; it returns a generator object. Each call to next() on that object runs the function until the next yield, hands out the value written there, and freezes with every variable intact until the following next().

def first_three():
    print("about to yield 1")
    yield 1
    print("about to yield 2")
    yield 2
    print("about to yield 3")
    yield 3

g = first_three()
print("created")
print(next(g))
print(next(g))

Output:

created
about to yield 1
1
about to yield 2
2

created printed first, before anything inside the function ran. The first next(g) ran lines 2 and 3 and stopped at the yield, so print received the 1. The second next resumed at line 4. Line 6 has not run and will not until a third next arrives. A next after the final yield raises StopIteration to signal that the generator has nothing more. try and except handle it like any other exception.

def two():
    yield "a"
    yield "b"

g = two()
print(next(g))
print(next(g))
try:
    print(next(g))
except StopIteration:
    print("generator is exhausted")

Output:

a
b
generator is exhausted

Looping over a generator

A for loop calls next() for you and stops quietly at StopIteration, which makes a generator a natural fit for for. The generator below wraps a file: it reads 1 line and yields it converted, then reads the next line only when the loop asks. int() ignores surrounding whitespace, so the \n ending each line causes no error.

import io

def plays(f):
    for line in f:
        yield int(line)

f = io.StringIO("3\n7\n5\n")
total = 0
for n in plays(f):
    total = total + n
    print("running total", total)

Output:

running total 3
running total 10
running total 15

At no point does the program hold more than 1 line of the file. The run below stands a list of 3 lines in for the file and doubles each value on its way through. Watch n and total change: each time the generator yields, its frame steps off the stack with line and its place in the loop intact, then steps back on when the loop asks for the next value. Python discards that frame only after the last line has been read.

def doubled(lines):
    for line in lines:
        yield int(line) * 2

lines = ["3\n", "7\n", "5\n"]
total = 0
for n in doubled(lines):
    total = total + n
print("total", total)

total 30

A generator runs once

A generator remembers how far it got, and when it reaches the end it stays there. A second loop over the same generator object receives nothing and its block never runs. To go through the data again, call the generator function again for a fresh object.

def nums():
    yield 1
    yield 2

g = nums()
for n in g:
    print("first pass", n)
for n in g:
    print("second pass", n)
print("done")

Output:

first pass 1
first pass 2
done

Practice

In what order do the printed lines appear?

def gen():
    print("start")
    yield "x"
    print("middle")
    yield "y"
    print("end")

g = gen()
a = next(g)
print(a)
b = next(g)
print(b)

• start, middle, end, x, y
• start, x, middle, y
• start, x, middle, y, end
• x, y

One yield written inside a for loop produces as many values as the loop has passes. Each request from the outer loop resumes the generator just after the yield, finishes that pass, takes the next line and pauses at the yield again. Here is plays(f) from earlier, 1 request at a time. Its input starts with 3 and 7 and ends with 5:

About this lesson

This section is part of the Python Foundations: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.