python - 如何编写生成器类？

Question

我看到了很多生成器函数的例子，但我想知道如何为类编写生成器。可以说，我想把斐波那契数列写成一门课。

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

输出:

<generator object __next__ at 0x000000000A3E4F68>
<generator object __next__ at 0x000000000A3E4F68>
<generator object __next__ at 0x000000000A3E4F68>

为什么值 self.a 没有被打印出来？另外，如何为生成器编写 unittest ？

Answer 1

How to write a generator class?

你快到了，正在编写一个 Iterator 类(我在答案末尾显示了一个生成器)，但是每次调用对象时都会调用 __next__使用 next，返回一个生成器对象。相反，为了让您的代码以最少的更改和最少的代码行运行，请使用 __iter__，这会使您的类实例化一个 iterable(这在技术上不是发电机):

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

当我们将一个可迭代对象传递给 iter() 时，它给了我们一个 iterator:

>>> f = iter(Fib())
>>> for i in range(3):
...     print(next(f))
...
0
1
1

要使类本身成为一个迭代器，它确实需要一个__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

现在，由于 iter 只是返回实例本身，我们不需要调用它:

>>> f = Fib()
>>> for i in range(3):
...     print(next(f))
...
0
1
1

Why is the value self.a not getting printed?

这是你的原始代码和我的评论:

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

所以每次调用 next(f) 时都会得到 __next__ 返回的生成器对象:

<generator object __next__ at 0x000000000A3E4F68>
<generator object __next__ at 0x000000000A3E4F68>
<generator object __next__ at 0x000000000A3E4F68>

Also, how do I write unittest for generators?

你仍然需要为 Generator

实现一个 send 和 throw 方法

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)

现在:

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

所以让我们实现一个生成器对象，并利用集合模块中的 Generator 抽象基类(参见其源代码 implementation )，这意味着我们只需要实现 send 和 throw - 给我们 close、__iter__(返回 self)和 __next__(与.send(None)) 免费(参见 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
    

并使用上述相同的测试:

>>> unittest.main(exit=False)
...
----------------------------------------------------------------------
Ran 3 tests in 0.002s

OK
<unittest.main.TestProgram object at 0x00000000031F7CC0>

Python 2

ABC Generator 只存在于 Python 3 中。要在没有 Generator 的情况下做到这一点，我们至少需要编写 close、 __iter__ 和 __next__ 除了我们上面定义的方法。

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")

请注意，我直接从 Python 3 standard library 复制了 close , 无需修改。