Python Generator Class

Ammar Ali Oct 10, 2023
  1. Python Iterators
  2. Python Generator Class
Python Generator Class

This tutorial will discuss the use of the yield statement and the next() function to create a generator class in Python.

To understand generators, we first need to understand the iterators discussed below.

Python Iterators

Iterators are objects used to access elements in a container one by one. We can use the for statement to loop over container objects to get the values individually.

An example code is shown below.

for element in [5, 6, 7]:
    print(element)

In the above Python code, we are looping over the list of elements and printing them one by one. Let’s understand what is happening behind the scenes.

The for statement calls the iter() function on the given container object, and the function contains a method __next__(), which will access each element of the given container object one by one.

The loop will terminate when the __next__() function raises an exception StopIteration, and the exception will only be raised when no more elements are present inside the given container object.

Python also provides the built-in function next() which can be used to call the __next__() function. To use the next() function on a container object, we have to create an object using the iter() function.

For example, let’s use a list of numbers and call the next() function to get each element of the list one by one. See the code and output below.

My_list = [5, 6, 7]
iter_object = iter(My_list)
print(next(iter_object))
print(next(iter_object))
print(next(iter_object))
print(next(iter_object))

Output:

5
6
7
---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-10-aa02bcda701b> in <module>
      4 print(next(iter_object))
      5 print(next(iter_object))
----> 6 print(next(iter_object))

StopIteration:

In the above code, we called the next() function four times, individually returning the three elements in the given list object. It returned the StopIteration exception when we called it the fourth time because no more elements were present in the list object.

We can also call the next() function using a loop.

Using the try-except statement, we can avoid the error and use the exception name to terminate the loop. For example, let’s repeat the above code using a loop and a try-except statement.

See the code and output below.

My_list = [5, 6, 7]
iter_object = iter(My_list)

for i in range(len(My_list)):
    try:
        print(next(iter_object))
    except StopIteration:
        break

Output:

5
6
7

In the above code, we used the name of the exception, StopIteration, to break the loop. The above iterator returns values one by one in a forward sequence, but we can also define our own iterator, which will return values according to our requirements.

We have to define three functions, __init__(), __iter__(), and __next__(), to add iterator behavior to a class. For example, let’s create a class that returns the Fibonacci numbers.

See the code below

class Fibexample:
    def __init__(self):
        self.x, self.y = 0, 1

    def __iter__(self):
        return self

    def __next__(self):
        r_value = self.x
        self.x, self.y = self.y, self.x + self.y
        return r_value


fib = Fibexample()
for i in range(7):
    print(next(fib))

Output:

0
1
1
2
3
5
8

The above class will return a number from the Fibonacci series whenever the next() function is called. In the above code, we called the next() function seven times, returning the first seven numbers of the Fibonacci series.

Python Generator Class

In Python, generators are used to create iterators. They are the same as regular functions; the only difference is using the yield statement instead of the return statement.

The yield() statement will call the next() function which returns an iterator object. For example, let’s create a generator function that returns the same Fibonacci series as the above code.

See the code and output below.

def fibexample(data_input):
    x, y = 0, 1
    for index in range(data_input):
        z = x
        x, y = y, x + y
        yield z


obj = fibexample(7)
for i in obj:
    print(i)

Output:

0
1
1
2
3
5
8

In the above code, the fibexample() function will return the required numbers of the Fibonacci series in an iterator object. We can use a loop to iterate through the object to get each value present in the iterator object.

The generator remembers the data values and the last execution of the next() function and will resume where it left off when the next() function is called again.

The result of the above function is the same as what we got in the iterators example, but the above code is relatively short compared to the code we used in the iterators example. The benefit of using the generators is that the __iter__() and __next__() functions will be created automatically, and the generators will also handle the StopIteration exception.

So, it’s easy to write iterators using the generators because creating iterators using generators is like writing a simple function using the yield statement.

Author: Ammar Ali
Ammar Ali avatar Ammar Ali avatar

Hello! I am Ammar Ali, a programmer here to learn from experience, people, and docs, and create interesting and useful programming content. I mostly create content about Python, Matlab, and Microcontrollers like Arduino and PIC.

LinkedIn Facebook

Related Article - Python Class