What is the Python Yield Statement?

Learn about the Python yield statement which is used to create generators.

“Count your age by friends, not years. Count your life by smiles, not tears.”
― John Lennon

1. Introduction

Python provides an yield statement which allows you to create generator functions. What is a generator function and how does the yield statement help with it? Let us find out in this article.

2. A Function vs a Generator

A function is defined in python using the def keyword. For example, the following function returns the square of the argument.

def joe(num):
    return num * num

Such a function can be invoked by passing it the required arguments, and it returns a value.

print joe(10)
# prints: 10

Now let us see what happens when the return keyword is replaced with yield.

def jack(num):
    yield num * num

On invoking this “function” in a similar manner, it returns a “generator object” as shown below.

print jack(10)
# prints <generator object jack at 0x327c550>

And that is how a generator is defined – any function which has an an yield statement within the body is a generator. On invocation, it does not execute the body, but returns a generator object as shown above.

3. What is a Generator?

So how do you run the function defined as a generator if it does not run when invoked normally? And what is this generator object that the function returns?

A generator is an enhanced iterator which provides the next() method, so it can be used in a for loop:

gn = jack(10)
for x in gn:
    print x
# prints: 100

Let us see what happens when you invoke the next() method directly.

xn = jack(10)
print xn.next()
# prints: 100

print xn.next()
# throws exception - StopIteration

In other words, the next() method of a generator can be executed repeatedly till it throws a StopIteration.

So, a generator is a function that can be used as the target of a for loop – it continues returning values as long as the yield statement is encountered. Once there are no more yield statements, the function raises a StopIteration exception.

It helps to think of a generator (a function containing an yield statement) as an iterator function.

4. Generating Values for a Loop

Let us now enhance our function to generate values for a loop. State is preserved and repeatedly calling the next() method continues execution where the last iteration left off.

def james(num):
    for i in xrange(num):
        yield i * i

Run the generator in a loop and see what happens:

v = james(10)
for x in v:
    print x
# prints:
0
1
4
9
16
25
36
49
64
81

So a generator can be used for dynamically generating values for a loop. Note that no storage of the values is involved.

5. Using return with yield

What happens when a function uses both the return statement, and the yield statement? Let us find out!

def james(num):
    for i in xrange(num):
        if i == 5: return
        else: yield i*i

Note: you cannot use the return statement with a value in a generator.

Again, running the loop, we find out that using the return statement inside a generator serves as a shortcut to end the loop. In other words, it works like a break statement in a loop.

v = james(10)
for x in v:
    print x
# prints
0
1
4
9
16

6. Restarting a Generator

You cannot restart a generator instance. One purpose of a generator is to not store the generated values, so restarting it means creating another instance of the generator.

7. Sending a Value into a Generator

Let us now look at some advanced features of a generator.

You can send a value into a generator function using the method send(). This value will be returned by the yield statement in the next round of iteration. Study the following example and the output carefully.

def joe():
    i = 5
    x = yield i*i
    print 'yield returned: ', x
    if x == 10: return
    
fn = joe()
print fn.next()
fn.send(10)
print fn.next()
# prints
25
yield returned:  10

What happens here is as follows:

  • A generator is created by invoking the function joe() and stored in a variable fn.
  • Printing the result of the first invocation shows: 25.
  • At this point, the yield statement has not yet returned in joe().
  • The caller sends a value of 10 into the generator.
  • The execution continues in joe() and returns the sent value as the result of yield.
  • The return statement is executed in joe() and the generator has completed.
  • The next time next() is invoked, it throws StopIteration since the generator is now complete.

8. On Demand Termination of an Iterator

This functionality of being able to send() values into a generator can be used to terminate a potentially infinite iterator on demand.

def joe():
    i = 0
    while True:
        x = yield i * i
        if x == 10: return
        else: i += 1

f = joe()
for x in f:
    print x
    if x > 20: f.send(10)

# prints
0
1
4
9
16
25

This provides a way to break out of the loop when some condition is fulfilled. However, that can also be done using a break statement. This method allows you to perform any cleanup necessary in the generator function.

Summary

To summarize, a generator is a function which contains an yield statement. It can be used to create an iterator without storing the intermediate values. Additionally being able to send values into a generator can be used for tweaking the generator while it is running.