How to write a generator class?

Question:

I see lot of examples of generator functions, but I want to know how to write generators for classes. Lets say, I wanted to write Fibonacci series as a class.

class Fib:
    def __init__(self):
        self.a, self.b = 0, 1

    def __next__(self):
        yield self.a
        self.a, self.b = self.b, self.a+self.b

f = Fib()

for i in range(3):
    print(next(f))

Output:

<generator object __next__ at 0x000000000A3E4F68>
<generator object __next__ at 0x000000000A3E4F68>
<generator object __next__ at 0x000000000A3E4F68>

Why is the value self.a not getting printed? Also, how do I write unittest for generators?

Asked By: Pritam

||

Answers:

__next__ should return an item, not yield it.

You can either write the following, in which Fib.__iter__ returns a suitable iterator:

class Fib:
    def __init__(self, n):
        self.n = n
        self.a, self.b = 0, 1

    def __iter__(self):
        for i in range(self.n):
            yield self.a
            self.a, self.b = self.b, self.a+self.b

f = Fib(10)

for i in f:
    print i

or make each instance itself an iterator by defining __next__.

class Fib:
    def __init__(self):
        self.a, self.b = 0, 1

    def __iter__(self):
        return self

    def __next__(self):
        x = self.a
        self.a, self.b = self.b, self.a + self.b
        return x

f = Fib()

for i in range(10):
    print next(f)
Answered By: chepner

How to write a generator class?

You’re almost there, writing an Iterator class (I show a Generator at the end of the answer), but __next__ gets called every time you call the object with next, returning a generator object. Instead, to make your code work with the least changes, and the fewest lines of code, use __iter__, which makes your class instantiate an iterable (which isn’t technically a generator):

class Fib:
    def __init__(self):
        self.a, self.b = 0, 1
    def __iter__(self):
        while True:
            yield self.a
            self.a, self.b = self.b, self.a+self.b

When we pass an iterable to iter(), it gives us an iterator:

>>> f = iter(Fib())
>>> for i in range(3):
...     print(next(f))
...
0
1
1

To make the class itself an iterator, it does require a __next__:

class Fib:
    def __init__(self):
        self.a, self.b = 0, 1        
    def __next__(self):
        return_value = self.a
        self.a, self.b = self.b, self.a+self.b
        return return_value
    def __iter__(self):
        return self

And now, since iter just returns the instance itself, we don’t need to call it:

>>> f = Fib()
>>> for i in range(3):
...     print(next(f))
...
0
1
1

Why is the value self.a not getting printed?

Here’s your original code with my comments:

class Fib:
    def __init__(self):
        self.a, self.b = 0, 1
        
    def __next__(self):
        yield self.a          # yield makes .__next__() return a generator!
        self.a, self.b = self.b, self.a+self.b

f = Fib()

for i in range(3):
    print(next(f))

So every time you called next(f) you got the generator object that __next__ returns:

<generator object __next__ at 0x000000000A3E4F68>
<generator object __next__ at 0x000000000A3E4F68>
<generator object __next__ at 0x000000000A3E4F68>

Also, how do I write unittest for generators?

You still need to implement a send and throw method for a Generator

from collections.abc import Iterator, Generator
import unittest

class Test(unittest.TestCase):
    def test_Fib(self):
        f = Fib()
        self.assertEqual(next(f), 0)
        self.assertEqual(next(f), 1)
        self.assertEqual(next(f), 1)
        self.assertEqual(next(f), 2) #etc...
    def test_Fib_is_iterator(self):
        f = Fib()
        self.assertIsInstance(f, Iterator)
    def test_Fib_is_generator(self):
        f = Fib()
        self.assertIsInstance(f, Generator)

And now:

>>> unittest.main(exit=False)
..F
======================================================================
FAIL: test_Fib_is_generator (__main__.Test)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "<stdin>", line 7, in test_Fib_is_generator
AssertionError: <__main__.Fib object at 0x00000000031A6320> is not an instance of <class 'collections.abc.Generator'>

----------------------------------------------------------------------
Ran 3 tests in 0.001s

FAILED (failures=1)
<unittest.main.TestProgram object at 0x0000000002CAC780>

So let’s implement a generator object, and leverage the Generator abstract base class from the collections module (see the source for its implementation), which means we only need to implement send and throw – giving us close, __iter__ (returns self), and __next__ (same as .send(None)) for free (see the Python data model on coroutines):

class Fib(Generator):
    def __init__(self):
        self.a, self.b = 0, 1        
    def send(self, ignored_arg):
        return_value = self.a
        self.a, self.b = self.b, self.a+self.b
        return return_value
    def throw(self, type=None, value=None, traceback=None):
        raise StopIteration
    

and using the same tests above:

>>> unittest.main(exit=False)
...
----------------------------------------------------------------------
Ran 3 tests in 0.002s

OK
<unittest.main.TestProgram object at 0x00000000031F7CC0>

Python 2

The ABC Generator is only in Python 3. To do this without Generator, we need to write at least close, __iter__, and __next__ in addition to the methods we defined above.

class Fib(object):
    def __init__(self):
        self.a, self.b = 0, 1        
    def send(self, ignored_arg):
        return_value = self.a
        self.a, self.b = self.b, self.a+self.b
        return return_value
    def throw(self, type=None, value=None, traceback=None):
        raise StopIteration
    def __iter__(self):
        return self
    def next(self):
        return self.send(None)
    def close(self):
        """Raise GeneratorExit inside generator.
        """
        try:
            self.throw(GeneratorExit)
        except (GeneratorExit, StopIteration):
            pass
        else:
            raise RuntimeError("generator ignored GeneratorExit")

Note that I copied close directly from the Python 3 standard library, without modification.

Do not use yield in __next__ function and implement next also for compatibility with python2.7+

Code

class Fib:
    def __init__(self):
        self.a, self.b = 0, 1
    def __next__(self):
        a = self.a
        self.a, self.b = self.b, self.a+self.b
        return a
    def next(self):
        return self.__next__()

If you give the class an __iter__() method implemented as a generator, "it will automatically return an iterator object (technically, a generator object)" when called, so that object’s __iter__() and __next__() methods will be the ones used.

Here’s what I mean:

class Fib:
    def __init__(self):
        self.a, self.b = 0, 1

    def __iter__(self):
        while True:
            value, self.a, self.b = self.a, self.b, self.a+self.b
            yield value

f = Fib()

for i, value in enumerate(f, 1):
    print(value)
    if i > 5:
        break

Output:

0
1
1
2
3
5
Answered By: martineau

Using yield in a method makes that method a generator, and calling that method returns a generator iterator. next() expects a generator iterator which implements __next__() and returns an item. That is why yielding in __next__() causes your generator class to output generator iterators when next() is called on it.

https://docs.python.org/3/glossary.html#term-generator

When implementing an interface, you need to define methods and map them to your class implementation. In this case the __next__() method needs to call through to the generator iterator.

class Fib:
    def __init__(self):
        self.a, self.b = 0, 1
        self.generator_iterator = self.generator()

    def __next__(self):
        return next(self.generator_iterator)

    def generator(self):
        while True:
            yield self.a
            self.a, self.b = self.b, self.a+self.b

f = Fib()

for i in range(3):
    print(next(f))
# 0
# 1
# 1
Answered By: Jared Deckard
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.