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?

Answer 1

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.