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?
Answers:
Thank you for visiting the Q&A section on Magenaut. Please note that all the answers may not help you solve the issue immediately. So please treat them as advisements. If you found the post helpful (or not), leave a comment & I’ll get back to you as soon as possible.
Method 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.
Method 2
__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)
Method 3
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
Method 4
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__()
Method 5
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
All methods was sourced from stackoverflow.com or stackexchange.com, is licensed under cc by-sa 2.5, cc by-sa 3.0 and cc by-sa 4.0