Generators and yield in Python

The last post ended with twenty lines of class to count down from three. Here is the same thing.

Three lines

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

print(list(countdown(3)))

for i in countdown(3):
    print(i)

It prints:

[3, 2, 1]
3
2
1

Any function with yield in it is a generator function. Calling it does not run the body — it returns a generator, and a generator is an iterator. Python writes __iter__ and __next__ for you.

Nothing runs until you ask

This is the part worth watching closely.

def noisy():
    print('  starting')
    yield 1
    print('  between')
    yield 2
    print('  finishing')

g = noisy()
print('created — the body has not run')
print('got', next(g))
print('got', next(g))
try:
    next(g)
except StopIteration:
    print('done')

It prints:

created — the body has not run
  starting
got 1
  between
got 2
  finishing
done

Read the order. starting appears after created, not before it. The body runs only when something calls next, and it stops again the moment it hits a yield.

return ends a function and throws away everything local to it. yield suspends the function with its variables intact, and the next next() picks up on the following line. When the body finally ends, Python raises StopIteration — the same signal the last post caught by hand.

Why laziness is worth anything

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('list:', sys.getsizeof(squares_list), 'bytes')
print('gen: ', sys.getsizeof(squares_gen), 'bytes')
print('same total:', sum(squares_list) == sum(squares_gen))

It prints:

list: 8448728 bytes
gen:  200 bytes
same total: True

The list holds a million numbers. The generator holds a paused function. Both sum to the same value, and only one of them had to fit in memory.

Round brackets instead of square brackets is the whole difference. That is a generator expression — a list comprehension that never builds the list.

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

print(sum(n * n for n in nums if n % 2))
print(sorted((n * n for n in nums), reverse=True)[:3])

It prints:

35
[36, 25, 16]

When the generator expression is the only argument, you can drop its brackets — that is why sum(n * n for n in nums) reads the way it does.

Infinite is now allowed

A list of every natural number is not a thing you can have. A generator of them is fine, because it only ever holds the one it is on.

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

out = []
for n in naturals():
    if n * n > 500:
        break
    out.append(n * n)

print(out)

It prints:

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484]

The while True never ends. The break is what ends it, and the generator is simply left paused, forever, and collected.

Stages that pull from each other

Generators chain. Each stage asks the one before it for a single item.

lines = ['17,ada,engineer', '', '18,grace,admiral', 'not a row', '19,alan,logician']

good   = (l for l in lines if l.count(',') == 2)
fields = (l.split(',') for l in good)
people = ((int(i), name) for i, name, _ in fields)

for pid, name in people:
    print(pid, name)

It prints:

17 ada
18 grace
19 alan

Nothing happened until the for loop asked. Then one line moved through all three stages, was printed, and the next one followed. At no point did a full intermediate list exist. Swap lines for a 4GB file and the code does not change.

yield from

When a generator’s job is to hand on what another one produces, yield from says it in one line.

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

print(list(flatten([1, [2, [3, 4]], 5, [[6]]])))

It prints:

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

Without it you would write a loop that re-yields each item, which works and reads worse.

The two things you give up

A generator is an iterator, so it runs once — everything the last post said applies here.

g = (n for n in range(4))
print(list(g))
print(list(g))

It prints:

[0, 1, 2, 3]
[]

And it has no length and no indexing, because it does not know what it has not produced yet.

g = (n for n in range(5))
try:
    print(len(g))
except TypeError as err:
    print('TypeError:', err)

try:
    print(g[0])
except TypeError as err:
    print('TypeError:', err)

It prints:

TypeError: object of type 'generator' has no len()
TypeError: 'generator' object is not subscriptable

If you need the length or an index, you needed a list. Call list() on it and take the memory.

What to remember

• yield in a function body makes it a generator function. Calling it runs nothing.

• Each next() runs the body up to the next yield and freezes it there, local variables and all.

• Round brackets give you a generator expression: the comprehension you already write, without the list.

• Generators can be infinite, and they chain into pipelines that never hold more than one item per stage.

• One pass, no len, no indexing. That is the price, and it is usually worth paying.