Generator Functions: yield Versus return

In this part of the book, we’ve learned about coding normal functions that receive input parameters and send back a single result immediately. It is also possible, however, to write functions that may send back a value and later be resumed, picking up where they left off. Such functions are known as generator functions because they generate a sequence of values over time.

Generator functions are like normal functions in most respects, and in fact are coded with normal def statements. However, when created, they are automatically made to implement the iteration protocol so that they can appear in iteration contexts. We studied iterators in Chapter 14; here, we’ll revisit them to see how they relate to generators.

广告：个人专属 VPN，独立 IP，无限流量，多机房切换，还可以屏蔽广告和恶意软件，每月最低仅 5 美元

State suspension

Unlike normal functions that return a value and exit, generator functions automatically suspend and resume their execution and state around the point of value generation. Because of that, they are often a useful alternative to both computing an entire series of values up front and manually saving and restoring state in classes. Because the state that generator functions retain when they are suspended includes their entire local scope, their local variables retain information and make it available when the functions are resumed.

The chief code difference between generator and normal functions is that a generator yields a value, rather than returning one—the yield statement suspends the function and sends a value back to the caller, but retains enough state to enable the function to resume from where it left off. When resumed, the function continues execution immediately after the last yield run. From the function’s perspective, this allows its code to produce a series of values over time, rather than computing them all at once and sending them back in something like a list.

Iteration protocol integration

To truly understand generator functions, you need to know that they are closely bound up with the notion of the iteration protocol in Python. As we’ve seen, iterable objects define a __next__ method, which either returns the next item in the iteration, or raises the special StopIteration exception to end the iteration. An object’s iterator is fetched with the iter built-in function.

Python for loops, and all other iteration contexts, use this iteration protocol to step through a sequence or value generator, if the protocol is supported; if not, iteration falls back on repeatedly indexing sequences instead.

To support this protocol, functions containing a yield statement are compiled specially as generators. When called, they return a generator object that supports the iteration interface with an automatically created method named __next__ to resume execution. Generator functions may also have a return statement that, along with falling off the end of the def block, simply terminates the generation of values—technically, by raising a StopIteration exception after any normal function exit actions. From the caller’s perspective, the generator’s __next__ method resumes the function and runs until either the next yield result is returned or a StopIteration is raised.

The net effect is that generator functions, coded as def statements containing yield statements, are automatically made to support the iteration protocol and thus may be used in any iteration context to produce results over time and on demand.

Note

As noted in Chapter 14, in Python 2.6 and earlier, iterable objects define a method named next instead of __next__. This includes the generator objects we are using here. In 3.0 this method is renamed to __next__. The next built-in function is provided as a convenience and portability tool: next(I) is the same as I.__next__() in 3.0 and I.next() in 2.6. Prior to 2.6, programs simply call I.next() instead to iterate manually.

Generator functions in action

To illustrate generator basics, let’s turn to some code. The following code defines a generator function that can be used to generate the squares of a series of numbers over time:

>>> def gensquares(N):
...     for i in range(N):
...         yield i ** 2        # Resume here later
...

This function yields a value, and so returns to its caller, each time through the loop; when it is resumed, its prior state is restored and control picks up again immediately after the yield statement. For example, when it’s used in the body of a for loop, control returns to the function after its yield statement each time through the loop:

>>> for i in gensquares(5):     # Resume the function
...     print(i, end=' : ')     # Print last yielded value
...
0 : 1 : 4 : 9 : 16 :
>>>

To end the generation of values, functions either use a return statement with no value or simply allow control to fall off the end of the function body.

If you want to see what is going on inside the for, call the generator function directly:

>>> x = gensquares(4)
>>> x
<generator object at 0x0086C378>

You get back a generator object that supports the iteration protocol we met in Chapter 14—the generator object has a __next__ method that starts the function, or resumes it from where it last yielded a value, and raises a StopIteration exception when the end of the series of values is reached. For convenience, the next(X) built-in calls an object’s X.__next__() method for us:

>>> next(x)                     # Same as x.__next__() in 3.0
0
>>> next(x)                     # Use x.next() or next() in 2.6
1
>>> next(x)
4
>>> next(x)
9
>>> next(x)

Traceback (most recent call last):
...more text omitted...
StopIteration

As we learned in Chapter 14, for loops (and other iteration contexts) work with generators in the same way—by calling the __next__ method repeatedly, until an exception is caught. If the object to be iterated over does not support this protocol, for loops instead use the indexing protocol to iterate.

Note that in this example, we could also simply build the list of yielded values all at once:

>>> def buildsquares(n):
...     res = []
...     for i in range(n): res.append(i ** 2)
...     return res
...
>>> for x in buildsquares(5): print(x, end=' : ')
...
0 : 1 : 4 : 9 : 16 :

For that matter, we could use any of the for loop, map, or list comprehension techniques:

>>> for x in [n ** 2 for n in range(5)]:
...     print(x, end=' : ')
...
0 : 1 : 4 : 9 : 16 :

>>> for x in map((lambda n: n ** 2), range(5)):
...     print(x, end=' : ')
...
0 : 1 : 4 : 9 : 16 :

However, generators can be better in terms of both memory use and performance. They allow functions to avoid doing all the work up front, which is especially useful when the result lists are large or when it takes a lot of computation to produce each value. Generators distribute the time required to produce the series of values among loop iterations.

Moreover, for more advanced uses, generators can provide a simpler alternative to manually saving the state between iterations in class objects—with generators, variables accessible in the function’s scopes are saved and restored automatically.[44] We’ll discuss class-based iterators in more detail in Part VI.

Extended generator function protocol: send versus next

In Python 2.5, a send method was added to the generator function protocol. The send method advances to the next item in the series of results, just like __next__, but also provides a way for the caller to communicate with the generator, to affect its operation.

Technically, yield is now an expression form that returns the item passed to send, not a statement (though it can be called either way—as yield X, or A = (yield X)). The expression must be enclosed in parentheses unless it’s the only item on the right side of the assignment statement. For example, X = yield Y is OK, as is X = (yield Y) + 42.

When this extra protocol is used, values are sent into a generator G by calling G.send(value). The generator’s code is then resumed, and the yield expression in the generator returns the value passed to send. If the regular G.__next__() method (or its next(G) equivalent) is called to advance, the yield simply returns None. For example:

>>> def gen():
...    for i in range(10):
...        X = yield i
...        print(X)
...
>>> G = gen()
>>> next(G)              # Must call next() first, to start generator
0
>>> G.send(77)           # Advance, and send value to yield expression
77
1
>>> G.send(88)
88
2
>>> next(G)              # next() and X.__next__() send None
None
3

The send method can be used, for example, to code a generator that its caller can terminate by sending a termination code, or redirect by passing a new position. In addition, generators in 2.5 also support a throw(type) method to raise an exception inside the generator at the latest yield, and a close method that raises a special GeneratorExit exception inside the generator to terminate the iteration. These are advanced features that we won’t delve into in more detail here; see reference texts and Python’s standard manuals for more information.

Note that while Python 3.0 provides a next(X) convenience built-in that calls the X.__next__() method of an object, other generator methods, like send, must be called as methods of generator objects directly (e.g., G.send(X)). This makes sense if you realize that these extra methods are implemented on built-in generator objects only, whereas the __next__ method applies to all iterable objects (both built-in types and user-defined classes).