user1022894 user1022894 -4 years ago 99
Python Question

How to write a generator class?

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?

Answer Source

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, use __iter__:

>>> 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
...
>>> f = iter(Fib())
>>> for i in range(3):
...     print(next(f))
...
0
1
1

To make the class itself an iterator:

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:

>>> f = Fib()
>>> f = iter(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 import Iterator, Generator
import unittest

class Test(unittest.TestCase):
    def test_Fib(self):
        f = Fib()
        self.assertEquals(next(f), 0)
        self.assertEquals(next(f), 1)
        self.assertEquals(next(f), 1)
        self.assertEquals(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:

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)
__main__:4: DeprecationWarning: Please use assertEqual instead.
...
----------------------------------------------------------------------
Ran 3 tests in 0.002s

OK
<unittest.main.TestProgram object at 0x00000000031F7CC0>
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download