python generators in a class

Question

I would like to create a class that has a method acting like a generator. I am pretty new to the generators. Here's my code:

class MyGen:

    def __init__(self):
        self.a = 0

    def create_generator(self):
        yield self.a
        self.a += 1

if __name__ == "__main__":
    myGenInstance = MyGen()
    myGen = myGenInstance.create_generator()
    for i in range(3):
        print(next(myGen))

I'm getting this error:

Traceback (most recent call last):
0
File "C:/Users/tvavr/PycharmProjects/filter/geneŕator_test.py", line 17, in <module>
print(next(myGen))
StopIteration
Process finished with exit code 1

What am I missing?

Answer 1

When reading your code, there is no error and the exception raised by python is a normal behavior.

the create_generator method can only yield one value because there was only one yield used.

consider that:

class Mygen:
    def __init__(self):
        self.a = 0

    def create_generator(self):
        yield self.a
        self.a += 1

        yield self.a
        self.a += 1

        yield self.a

When you run the code as follow notice what happend:

if __name__ == "__main__":
my_gen_instance = MyGen()
my_gen = my_gen_instance.create_generator()
for i in range(3):
    print('the value number', i, 'is:', next(myGen))

the output will be:

the value number 0 is 0
the value number 1 is 1
the value number 2 is 2

Now if we change the range from 3 to 4, see what happens:

if __name__ == "__main__":
    my_gen_instance = MyGen()
    my_gen = my_gen_instance.create_generator()
    for i in range(4):
        print('the value number', i, 'is:', next(myGen))

the output will be:

the value number 0 is 0
the value number 1 is 1
the value number 2 is 2
Traceback (most recent call last):
  File "pygen.py", line 21, in <module>
    print('the value number', i, 'is', next(my_gen))
StopIteration

Because generator are special type of iterator they throw a StopIteration when they are exhausted like normal iterator

So to avoid that you have to anticipate this behavior by including as many yields as you need and the best way is to create an infinite generator inside the MyGen class as follows:

class MyGen:
    def __init__(self):
        self.a = 0

    def create_generator(self):
        while True:
            yield self.a
            self.a += 1

Now you can iterate as much as you want over the generator created by this method.