问题
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?
回答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, 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 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.
回答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)
回答3:
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__()
回答4:
If you give the class an __iter__() method implemented as a generator, it will automatically return 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
来源：https://stackoverflow.com/questions/42983569/how-to-write-a-generator-class