6. Generators

Generator functions are one of Python’s most interesting and powerful features. Generators are often presented as a convenient way to define new kinds of iteration patterns. However, there is much more to them: Generators can also fundamentally change the whole execution model of functions. This chapter discusses generators, generator delegation, generator-based coroutines, and common applications of generators.

6.1 Generators and yield

If a function uses the yield keyword, it defines an object known as a generator. The primary use of a generator is to produce values for use in iteration. Here’s an example:

def countdown(n):
     print('Counting down from', n)
     while n > 0:
          yield n
          n -= 1

# Example use
for x in countdown(10):
    print('T-minus', x)

If you call this function, you will find that none of its code starts executing. For example:

>>> c = countdown(10)
>>> c
<generator object countdown at 0x105f73740
> >>>

Instead, a generator object is created. The generator object, in turn, only executes the function when you start iterating on it. One way to do that is to call next() on it:

>>> next(c)
Counting down from 10
10
>>> next(c)
 9

When next() is called, the generator function executes statements until it reaches a yield statement. The yield statement returns a result, at which point execution of the function is suspended until next() is invoked again. While it’s suspended, the function retains all of its local variables and execution environment. When resumed, execution continues with the statement following the yield.

next() is a shorthand for invoking the __next__() method on a generator. For example, you could also do this:

>>> c.__next__()
8
>>> c.__next__()
7
>>>

You normally don’t call next() on a generator directly, but use the for statement or some other operation that consumes the items. For example:

for n in countdown(10):
     statements

a = sum(countdown(10))

A generator function produces items until it returns—by reaching the end of the function or by using a return statement. This raises a StopIteration exception that terminates a for loop. If a generator function returns a non-None value, it is attached to the StopIteration exception. For example, this generator function uses both yield and return:

def func():
    yield 37
    return 42

Here’s how the code would execute:

>>> f = func()
>>> f
<generator object func at 0x10b7cd480>
>>> next(f)
37
>>> next(f)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration: 42
>>>

Observe that the return value is attached to StopIteration. To collect this value, you need to explicitly catch StopIteration and extract the value:

try:
    next(f)
except StopIteration as e:
    value = e.value

Normally, generator functions don’t return a value. Generators are almost always consumed by a for loop where there is no way to obtain the exception value. This means the only practical way to get the value is to drive the generator manually with explicit next() calls. Most code involving generators just doesn’t do that.

A subtle issue with generators is where a generator function is only partially consumed. For example, consider this code that abandons a loop early:

for n in countdown(10):
     if n == 2:
         break
     statements

In this example, the for loop aborts by calling break and the associated generator never runs to full completion. If it’s important for your generator function to perform some kind of cleanup action, make sure you use try-finally or a context manager. For example:

def countdown(n):
    print('Counting down from', n)
    try:
         while n > 0:
              yield n
              n = n - 1
    finally:
         print('Only made it to', n)

Generators are guaranteed to execute the finally block code even if the generator is not fully consumed—it will execute when the abandoned generator is garbage-collected. Similarly, any cleanup code involving a context manager is also guaranteed to execute when a generator terminates:

def func(filename):
    with open(filename) as file:
         ...
         yield data
         ...
    # file closed here even if generator is abandoned

Proper cleanup of resources is a tricky problem. As long as you use constructs such as try-finally or context managers, generators are guaranteed to do the right thing even if they are terminated early.

6.2 Restartable Generators

Normally a generator function executes only once. For example:

>>> c = countdown(3)
>>> for n in c:
...    print('T-minus', n)
...
T-minus 3
T-minus 2
T-minus 1
>>> for n in c:
...    print('T-minus', n)
...
>>>

If you want an object that allows repeated iteration, define it as a class and make the __iter__() method a generator:

class countdown:
    def __init__(self, start):
        self.start = start

    def __iter__(self):
        n = self.start
        while n > 0:
            yield n
            n -= 1

This works because each time you iterate, a fresh generator is created by __iter__().

6.3 Generator Delegation

An essential feature of generators is that a function involving yield never executes by itself—it always has to be driven by some other code using a for loop or explicit next() calls. This makes it somewhat difficult to write library functions involving yield because calling a generator function is not enough to make it execute. To address this, the yield from statement can be used. For example:

def countup(stop):
    n = 1
    while n <= stop:
        yield n
        n += 1

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

def up_and_down(n):
    yield from countup(n)
    yield from countdown(n)

yield from effectively delegates the iteration process to an outer iteration. For example, you would write code like this to drive the iteration:

>>> for x in up_and_down(5):
...    print(x, end=' ')
1 2 3 4 5 5 4 3 2 1
>>>

yield from mainly saves you from having to drive iteration yourself. Without this feature, you would have to write up_and_down(n) as follows:

def up_and_down(n):
    for x in countup(n):
        yield x
    for x in countdown(n):
        yield x

yield from is especially useful when writing code that must recursively iterate through nested iterables. For example, this code flattens nested lists:

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

Here is an example of how this works:

>>> a = [1, 2, [3, [4, 5], 6, 7], 8]
>>> for x in flatten(a):
...    print(x, end=' ')
...
1 2 3 4 5 6 7 8
>>>

One limitation of this implementation is that it is still subject to Python’s recursion limit, so it would not be able to handle deeply nested structures. This will be addressed in the next section.

6.4 Using Generators in Practice

At first glance, it might not be obvious how to use generators for practical problems beyond defining simple iterators. However, generators are particularly effective at structuring various data handling problems related to pipelines and workflows.

One useful application of generators is as a tool for restructuring code that consists of deeply nested for loops and conditionals. Consider this script that searches a directory of Python files for all comments containing the word “spam”:

import pathlib
import re


for path in pathlib.Path('.').rglob('*.py'):
    if path.exists():
        with path.open('rt', encoding='latin-1') as file:
            for line in file:
                m = re.match('.*(#.*)$', line)
                if m:
                     comment = m.group(1)
                     if 'spam' in comment:
                         print(comment)

Notice the number of levels of nested control flow. Your eyes are already starting to hurt as you look at the code. Now, consider this version using generators:

import pathlib
import re

def get_paths(topdir, pattern):
    for path in pathlib.Path(topdir).rglob(pattern)
        if path.exists():
            yield path

def get_files(paths):
    for path in paths:
        with path.open('rt', encoding='latin-1') as file:
             yield file

def get_lines(files):
    for file in files:
        yield from file

def get_comments(lines):
    for line in lines:
        m = re.match('.*(#.*)$', line)
        if m:
            yield m.group(1)

def print_matching(lines, substring):
    for line in lines:
        if substring in lines:
            print(substring)

paths = get_paths('.', '*.py')
files = get_files(pypaths)
lines = get_lines(pyfiles)
comments = get_comments(lines)
print_matching(comments, 'spam')

In this section, the problem is broken down into smaller self-contained components. Each component only concerns itself with a specific task. For example, the get_paths() generator is only concerned with path names, the get_files() generator is only concerned with opening files, and so forth. It is only at the end that these generators are hooked together into a workflow to solve a problem.

Making each component small and isolated is a good abstraction technique. For example, consider the get_comments() generator. As input, it takes any iterable producing lines of text. This text could come from almost anywhere—a file, a list, a generator, and so on. As a result, this functionality is much more powerful and adaptable than it was when it was embedded into a deeply nested for loop involving files. Generators thus encourage code reuse by breaking problems into small well-defined computational tasks. Smaller tasks are also easier to reason about, debug, and test.

Generators are also useful for altering the normal evaluation rules of function application. Normally, when you apply a function, it executes immediately, producing a result. Generators don’t do that. When a generator function is applied, its execution is delayed until some other bit of code invokes next() on it (either explicitly or by a for loop).

As an example, consider again the generator function for flattening nested lists:

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

One problem with this implementation is that, due to Python’s recursion limit, it won’t work with deeply nested structures. This can be fixed by driving iteration in a different way using a stack. Consider this version:

def flatten(items):
    stack = [ iter(items) ]
    while stack:
        try:
            item = next(stack[-1])
            if isinstance(item, list):
                stack.append(iter(item))
            else:
                yield item
        except StopIteration:
            stack.pop()

This implementation builds an internal stack of iterators. It is not subject to Python’s recursion limit because it’s putting data on an internal list as opposed to building frames on the internal interpreter stack. Thus, if you find yourself in need of flattening a few million layers of some uncommonly deep data structure, you’ll find that this version works fine.

Do these examples mean that you should rewrite all of your code using wild generator patterns? No. The main point is that the delayed evaluation of generators allows you to alter the spacetime dimensions of normal function evaluation. There are various real-world scenarios where these techniques can be useful and applied in unexpected ways.

6.5 Enhanced Generators and yield Expressions

Inside a generator function, the yield statement can also be used as an expression that appears on the right-hand side of an assignment operator. For example:

def receiver():
    print('Ready to receive')
    while True:
          n = yield
          print('Got', n)

A function that uses yield in this manner is sometimes known as an “enhanced generator” or “generator-based coroutine.” Sadly, this terminology is a bit imprecise and made even more confusing since “coroutines” are more recently associated with async functions. To avoid this confusion, we’ll use the term “enhanced generator” to make it clear that we’re still talking about standard functions that use yield.

A function that uses yield as an expression is still a generator, but its usage is different. Instead of producing values, it executes in response to values sent to it. For example:

>>> r = receiver()
>>> r.send(None)        # Advances to the first yield
Ready to receive
>>> r.send(1)
Got 1
>>> r.send(2)
Got 2
>>> r.send('Hello')
Got Hello
>>>

In this example, the initial call to r.send(None) is necessary so that the generator executes statements leading to the first yield expression. At this point, the generator suspends, waiting for a value to be sent to it using the send() method of the associated generator object r. The value passed to send() is returned by the yield expression in the generator. Upon receiving a value, a generator executes statements until the next yield is encountered.

As written, the function runs indefinitely. The close() method can be used to shut down the generator as follows:

>>> r.close()
>>> r.send(4)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>>

The close() operation raises a GeneratorExit exception inside the generator at the current yield. Normally, this causes the generator to terminate silently; if you’re so inclined, you can catch it to perform cleanup actions. Once closed, a StopIteration exception will be raised if further values are sent to a generator.

Exceptions can be raised inside a generator using the throw(ty [,val [,tb]]) method where ty is the exception type, val is the exception argument (or tuple of arguments), and tb is an optional traceback. For example:

>>> r = receiver()
Ready to receive
>>> r.throw(RuntimeError,  "Dead")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "receiver.py", line 14, in receiver
    n = yield
RuntimeError: Dead
>>>

Exceptions raised in either manner will propagate from the currently executing yield statement in the generator. A generator can elect to catch the exception and handle it as is appropriate. If a generator doesn’t handle the exception, it propagates out of the generator to be handled at a higher level.

6.6 Applications of Enhanced Generators

Enhanced generators are an odd programming construct. Unlike a simple generator which naturally feeds a for loop, there is no core language feature that drives an enhanced generator. Why, then, would you ever want a function that needs values to be sent to it? Is it purely academic?

Historically, enhanced generators have been used in the context of concurrency libraries—especially those based on asynchronous I/O. In that context, they’re usually referred to as “coroutines” or “generator-based coroutines.” However, much of that functionality has been folded into the async and await features of Python. There is little practical reason to use yield for that specific use case. That said, there are still some practical applications.

Like generators, an enhanced generator can be used to implement different kinds of evaluation and control flow. One example is the @contextmanager decorator found in the contextlib module. For example:

from contextlib import contextmanager

@contextmanager
def manager():
    print("Entering")
    try:
         yield 'somevalue'
    except Exception as e:
         print("An error occurred", e)
    finally:
         print("Leaving")

Here, a generator is being used to glue together the two halves of a context manager. Recall that context managers are defined by objects implementing the following protocol:

class Manager:
    def __enter__(self):
        return somevalue

    def __exit__(self, ty, val, tb):
        if ty:
            # An exception occurred
            ...
            # Return True/ if handled. False otherwise

With the @contextmanager generator, everything prior to the yield statement executes when the manager enters (via the __enter__() method). Everything after the yield executes when the manager exits (via the __exit__() method). If an error took place, it is reported as an exception on the yield statement. Here is an example:

>>> with manager() as val:
...        print(val)
...
Entering
somevalue
Leaving
>>> with manager() as val:
...        print(int(val))
...
Entering
An error occurred invalid literal for int() with base 10: 'somevalue'
Leaving
>>>

To implement this, a wrapper class is used. This is a simplified implementation that illustrates the basic idea:

class Manager:
    def __init__(self, gen):
        self.gen = gen

    def __enter__(self):
        # Run to the yield
        return self.gen.send(None)

    def __exit__(self, ty, val, tb):
        # Propagate an exception (if any)
        try:
            if ty:
                try:
                    self.gen.throw(ty, val, tb)
                except ty:
                    return False
            else:
                        self.gen.send(None)
                except StopIteration:
                    return True

Another application of extended generators is using functions to encapsulate a “worker” task. One of the central features of a function call is that it sets up an environment of local variables. Access to those variables is highly optimized—it’s much faster than accessing attributes of classes and instances. Since generators stay alive until explicitly closed or destroyed, one might use a generator to set up a long-lived task. Here is an example of a generator that receives byte fragments and assembles them into lines:

@consumer
def line_receiver():
    data = bytearray()
    line = None
    linecount = 0
    while True:
        part = yield line
        linecount += part.count(b'\n')
        data.extend(part)
        if linecount > 0:
            index = data.index(b'\n')
            line = bytes(data[:index+1])
            data = data[index+1:]
            linecount -= 1
        else:
            line = None

In this example, a generator receives byte fragments that are collected into a byte array. If the array contains a newline, a line is extracted and returned. Otherwise, None is returned. Here’s an example illustrating how it works:

>>> r = line_receiver()
>>> r.send(b'hello')
>>> r.send(b'world\nit ')
b'hello world\n'
>>> r.send(b'works!')
>>> r.send(b'\n')
b'it works!\n''
>>>

Similar code could be written as a class such as this:

class LineReceiver:
    def __init__(self):
        self.data = bytearray()
        self.linecount = 0

    def send(self, part):
        self.linecount += part.count(b'\n')
        self.data.extend(part)
        if self.linecount > 0:
            index = self.data.index(b'\n')
            line = bytes(self.data[:index+1])
            self.data = self.data[index+1:]
            self.linecount -= 1
            return line
        else:
            return None

Although writing a class might be more familiar, the code is more complex and runs slower. Tested on the author’s machine, feeding a large collection of chunks into a receiver is about 40–50% faster with a generator than with this class code. Most of those savings are due to the elimination of instance attribute lookup—local variables are faster.

Although there are many other potential applications, the important thing to keep in mind is that if you see yield being used in a context that is not involving iteration, it is probably using the enhanced features such as send() or throw().

6.7 Generators and the Bridge to Awaiting

A classic use of generator functions is in libraries related to asynchronous I/O such as in the standard asyncio module. However, since Python 3.5 much of this functionality has been moved into a different language feature related to async functions and the await statement (see the last part of Chapter 5).

The await statement involves interacting with a generator in disguise. Here is an example that illustrates the underlying protocol used by await:

class Awaitable:
    def __await__(self):
        print('About to await')
        yield    # Must be a generator
        print('Resuming')

# Function compatible with "await". Returns an "awaitable"
def function():
    return Awaitable()

async def main():
    await function()

Here is how you might try the code using asyncio:

>>> import asyncio
>>> asyncio.run(main())
About to await
Resuming
>>>

Is it absolutely essential to know how this works? Probably not. All of this machinery is normally hidden from view. However, if you ever find yourself using async functions, just know that there is a generator function buried somewhere inside. You’ll eventually find it if you just keep on digging the hole of technical debt deep enough.

6.8 Final Words: A Brief History of Generators and Looking Forward

Generators are one of Python’s more interesting success stories. They are also part of a greater story concerning iteration. Iteration is one of the most common programming tasks of all. In early versions of Python, iteration was implemented via sequence indexing and the __getitem__() method. This later evolved into the current iteration protocol based on __iter__() and __next__() methods. Generators appeared shortly thereafter as a more convenient way to implement an iterator. In modern Python, there is almost no reason to ever implement an iterator using anything other than a generator. Even on iterable objects that you might define yourself, the __iter__() method itself is conveniently implemented in this way.

In later versions of Python, generators took on a new role as they evolved enhanced features related to coroutines—the send() and throw() methods. These were no longer limited to iteration but opened up possibilities for using generators in other contexts. Most notably, this formed the basis of many so-called “async” frameworks used for network programming and concurrency. However, as asynchronous programming has evolved, most of this has transformed into later features that use the async/await syntax. Thus, it’s not so common to see generator functions used outside of the context of iteration—their original purpose. In fact, if you find yourself defining a generator function and you’re not performing iteration, you should probably reconsider your approach. There may be a better or more modern way to accomplish what you’re doing.