How to write a generator class?

Question

I see many examples of generator functions, but I want to know how to write generators for classes. Suppose I wanted to write a 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 printed? Also, how to write unittest for generators?

Answer 1

How to write a generator class?

You almost write the Iterator class (at the end of the answer I show Generator), but __next__ is called every time you call the object with next , returning the generator object. 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 = iter(Fib()) >>> for i in range(3): ... print(next(f)) ... 0 1 1 

Why is the value self.a not printed?

Here is 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)) 

Therefore, each time you call next(f) , you get a generator object that returns __next__ :

 <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 the send and throw method for 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 the generator object and use the abstract Generator base class from the collections module (see the source for its implementation ), which means that we only need to implement send and throw - it provides us with free close , __iter__ (returns self) and __next__ (similarly .send(None) ) (see Python data model for 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 as above:

 >>> unittest.main(exit=False) ... ---------------------------------------------------------------------- Ran 3 tests in 0.002s OK <unittest.main.TestProgram object at 0x00000000031F7CC0> 

Python 2

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 that 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 , unchanged.

Answer 2

__next__ should return an element, not give 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 an iterator itself 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) 

Answer 3

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

the 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__() 

Answer 4

If you give the class the __iter__() method implemented as a generator , it will automatically return the generator object when called, so that the __iter__ and __next__ will be used by those that are used.

Here is 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