Python's yield Keyword - How it Works

To return a value from a Python function, you must have used the return keyword. When the program encounters the return keyword, it exits the function and returns the value to the caller before proceeding to execute additional code.

However, the case is entirely different with the yield keyword. It does control the flow of execution and returns a value like the return keyword but how it behaves is different.

yield Keyword

When a function contains a yield statement, it is considered a generator function. Theyieldkeyword alters the behavior of the function. How?

The yield keyword produces a series of values and returns when requested by the caller.

The yield keyword temporarily freezes the execution of the function and returns the value to the caller. As it freezes the execution, the state of the function is preserved and allows the function to resume from where it was left off.

Working of yield when a function encounters it

yield in a Function

You can define a normal function as usual with the def keyword, but instead of the return keyword, use the yield keyword in the function body.

# A function with yield kw - a generator function
def random_func():
    # Yield values
    yield "random1"
    yield "random2"
    yield "random3"

# Instance of the function
obj = random_func()

The above code defines a function named random_func(), which is essentially a generator function due to the yield keyword in the function body.

Now, this is a simple generator function that yields three values and returns them when requested. The function is initialized (random_func()) and saved in the obj variable.

When you print the obj variable, you will receive the generator object.

...
print(obj)
--------------------
<generator object random_func at 0x0000026E95924880>

So, how do you access the yielded values one at a time? The values can be accessed by using the caller's __next__() method, as shown below.

# Instance of the function
obj = random_func()
# Calling __next__() method on the instance
print(obj.__next__())   # or print(next(obj))
print(obj.__next__())   # or print(next(obj))
print(obj.__next__())   # or print(next(obj))
--------------------
random1
random2
random3

The __next__() method is called three times to retrieve all three values returned by the function. You can also iterate over the obj variable with the for loop to get the values, which is more efficient than repeatedly calling the __next__() method when you need to consume all of the yielded values.

# Instance of the function
obj = random_func()
# Iterating through the instance
for value in obj:
    print(value)
--------------------
random1
random2
random3

How yield Works in a Function?

How does the execution flow in the function containing theyieldstatement? Let's keep it very simple and understand with an example.

import logging

# Configuring logging
logging.basicConfig(level=logging.INFO)

# Simple generator function
def generate_val():
    logging.info("Start")
    value = 0
    logging.info("Before entering the loop")
    while True:
        logging.info("Before yield statement")
        yield value
        logging.info("After yield statement")
        value += 1

gen = generate_val()
print(next(gen))
print(next(gen))
print(next(gen))

The above code defines a basic generator function generate_val() that generates an infinite sequence of values starting from 0. It uses the yield statement to yield each value and increments it by 1 in each iteration of the loop.

The logging module is used to log the information. It logs various messages at different points in the generator function's execution to provide information about its flow.

When you run the code, the following output will be generated.

0
1
2
INFO:root:Start
INFO:root:Before entering the loop
INFO:root:Before yield statement
INFO:root:After yield statement
INFO:root:Before yield statement
INFO:root:After yield statement
INFO:root:Before yield statement

Notice that when the program enters the loop and encounters the yield statement, it returns the value for the first iteration (next(gen)), suspends the execution flow immediately, and then resumes from where it left off for the next iteration.

This will happen with each iteration, and the program will never exit the function until stopped manually.

StopIteration Exception

In the above section, the program runs infinitely and never exhausts. What happens when the generator function has nothing else to evaluate?

Let's understand with a basic example.

# Function to yield only even number
def only_even(n):
    try:
        num = 0
        while num <= n:
            if num % 2 == 0:
                yield num
            num += 1
    except Exception as e:
        print(f"Error - {str(e)}.")

The code above defines the only_even() function, which accepts an argument n and returns even numbers up to n, including n, if it is an even number.

# Function to yield only even number
...

even = only_even(3)

The function (only_even()) is instantiated with the argument 3 and stored in the even variable.

The function can now generate even numbers up to 3 (0 and 2), which means you can use the next() method on "even" twice. What happens if you call the next() method more than twice?

# Function to yield only even number
...

even = only_even(3)
print(next(even))
print(next(even))
print(next(even))    # Raises an error

For the first two iterations, the function returns an even number. In the third iteration, the loop is executed until num <= n becomes false. If the generator produces no more items, the StopIteration error is returned.

0
2
Traceback (most recent call last):
  ...
    print(next(even))
          ^^^^^^^^^^
StopIteration

To avoid the StopIteration error, you can use a for loop. A for loop automatically catches StopIteration exceptions and terminates gracefully when the iterator is exhausted.

# Function to yield only even number
...

even = only_even(9)
for num in even:
    print(num)

--------------------
0
2

When to Use yield?

You can use a yield statement in a function when you want to generate a sequence of values lazily (generate values when requested) rather than computing and storing them in memory until you print them.

Let's say you want to create a function that returns all the even numbers between 0 and 100000000.

Using the return keyword in the function computes all of the values and stores them in memory, when you print them, the program dumps them all at once.

def only_even(n):
    number = []
    for num in range(n):
        if num % 2 == 0:
            number.append(num)
    return number

even = only_even(100000000)
for number in even:
    print(number)

If you include a yield statement in the function, you can generate values as they are needed.

def only_even(n):
    num = 0
    while num <= n:
        if num % 2 == 0:
            yield num
        num += 1

even = only_even(100000000)
for val in even:
    print(val)

So, what is the main difference between the two programs? If you run the first program (function with return keyword), the function will compute the values and store them in memory until the program reaches the print statement, at which point it will print them, whereas the second program (function with yield keyword) generates values on-the-fly as they are requested, rather than computing and storing all of the values in memory at once.

Conclusion

The yield keyword in a function body means that the function is a generator function and the yield produces a sequence of values and returns when the value is requested by the caller.

The yield keyword remembers the state of the function and allows the function to resume from where it was left followed by the first iteration.

Let's recall what you've learned: