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?
Advertisement
Answer
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.