Basics of Python Generator Functions

Python generator functions allow you to declare a function that behaves like an iterator, making it a faster, cleaner and easier way to create an iterator. An iterator is an object that can be iterated or looped upon. It is used to abstract a container of data to make it behave like an iterable object. Examples of iterable objects that are used more commonly include lists, dictionaries and strings.

In this article, we will learn to create and use generators in Python with the help of some examples.

Python Class Iterator Implementation

Let’s first look at a simple, class-based iterator code to produce odd numbers:

class get_odds:
    def __init__(self, max):
        self.n=3
        self.max=max
    def __iter__(self):
        return self
    def __next__(self):
        if self.n <= self.max:
            result = self.n
            self.n += 2
            return result
        else:
            raise StopIteration
numbers = get_odds(10)
print(next(numbers))
print(next(numbers))
print(next(numbers))
# Output
3
5
7

As you can see, a sequence of odd numbers are generated. To generate this, we created a custom iterator inside the get_odds class. For an object to be an iterator it should implement the __iter__ method, which will return the iterator object. The __next__ method will then return the next value in the sequence and might possibly raise the StopIteration exception when there are no values to be returned. As you can see, the process of creating iterators is lengthy, which is why we turn to generators. Again, Python generators are a simple way of implementing iterators.

More on Python: 13 Python Code Snippets You Need to Know

Python Generator Implementation 

Let’s use the previous code and implement the same iterator except using a Python generator.

def get_odds_generator():
    n=1
    
    n+=2
    yield n
    
    n+=2
    yield n 
    
    n+=2
    yield n
    
numbers=get_odds_generator()
print(next(numbers))
print(next(numbers))
print(next(numbers))
# Output
3
5
7

I first created a generator function that has three yield statements, and when we call this function, it returns a generator that is an iterator object. We then called the next() method to retrieve elements from this object. The first print statement gives us the value of the first yield, which is three; the second print statement gives us the value of the second yield statement, which is five; and the last print statement gives us the value of the third yield statement, which is seven. As you can see, the generator function is much simpler compared to our class-based iterator.

Now, let’s try to implement a loop to make this Python generator return odd numbers until a certain max number.

def get_odds_generator(max):
    n=1
    
    while n<=max:
        yield n
        n+=2
    
numbers=get_odds_generator(3)
print(next(numbers))
print(next(numbers))
print(next(numbers))

As you can see from the output, one and three were generated and after that, a StopIteration exception has been raised. The loop condition (n<=max) is False since max is three and n is five, therefore the StopIteration exception was raised.

When comparing this code with our get_odds class, you can see that in our generator, we never explicitly defined the __iter__ method, the __next__ method or raised a StopIteration exception. These are handled implicitly by generators, making programming much easier and simpler to understand.

Iterators and generators are typically used to handle a large stream of data, and theoretically, even an infinite stream of data. These large streams of data cannot be stored in memory at once. To handle this, we can use generators to handle only one item at a time. Next, we will build a generator to produce an infinite stream of Fibonacci numbers. Fibonacci numbers are a series of numbers where the next element is the sum of the previous two elements.

def fibonacci_generator():
    n1=0
    n2=1
    while True:
        yield n1
        n1, n2 = n2, n1 + n2
sequence= fibonacci_generator()
print(next(sequence))
print(next(sequence))
print(next(sequence))
print(next(sequence))
print(next(sequence))
# Output
0
1
1
2
3

In defining the fibonacci_generator function, I first created the first two elements of the Fibonacci series, then used an infinite while loop, and inside it, yield the value of n1, and then updated the values so that the next term will be the sum of the previous two terms with the line n1,n2=n2,n1+n2. Our print statements gave us the sequence of numbers in the Fibonacci sequence. If we had used a for loop and a list to store this infinite series, we would have run out of memory. However, with generators, we can keep accessing these terms for as long as we want since we are dealing with one item at a time.

More on Python: Introduction to Priority Queues in Python

Difference Between Python Generator Functions and Regular Functions

The main difference between a regular function and generator functions is that the state of generator functions is maintained through the use of the keyword yield and works much like using return. But there are some other important differences. 

For starters, yield saves the state of the function. The next time the function is called, execution continues from where it left off, with the same variable values it had before yielding, whereas the return statement terminates the function completely. Another difference is that generator functions don’t even run a function, it only creates and returns a generator object. Lastly, the code in generator functions only execute when next() is called on the generator object.

From this article, we have covered the basics of python generators. You can also create generators on the fly using generator expressions.