Как написать класс генератора?

Я вижу много примеров функций генератора, но я хочу знать, как писать генераторы для классов. Допустим, я хотел написать серию Фибоначчи как класс.

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 для генераторов?

4 ответа

Решение

Как написать класс генератора?

Вы почти там, пишете класс Iterator (я показываю Generator в конце ответа), но __next__ вызывается каждый раз, когда вы вызываете объект с next, возвращая объект генератора. Вместо этого используйте __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

Чтобы сделать сам класс итератором:

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

И сейчас:

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

Почему значение self.a не печатается?

Вот ваш оригинальный код с моими комментариями:

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>

Кроме того, как мне написать unittest для генераторов?

Вам все еще нужно реализовать метод send и throw для Generator

from collections 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 абстрактный базовый класс из модуля коллекций (см. источник для его реализации), что означает, что нам нужно только реализовать send а также throw - давая нам close, __iter__ (возвращает себя) и __next__ (такой же как .send(None)) бесплатно (см. модель данных Python на сопрограммах):

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

Азбука 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")

Обратите внимание, что я скопировал close непосредственно из стандартной библиотеки Python 3, без изменений.

__next__ должен вернуть товар, а не отдать его.

Вы можете написать следующее, в котором Fib.__iter__ возвращает подходящий итератор:

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

или сделать каждый экземпляр самим итератором, определив __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)

Если вы дадите классу __iter__() Метод реализован как генератор, он будет автоматически возвращать объект генератора при вызове, так что объект __iter__ а также __next__ методы будут те, которые используются.

Вот что я имею в виду:

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

Выход:

0
1
1
2
3
5

С помощью yieldв методе делает этот метод генератором, а вызов этого метода возвращает итератор генератора.next() ожидает итератор генератора, который реализует __next__() а также returnэто предмет. Поэтомуyieldвходить в __next__() заставляет ваш класс генератора выводить итераторы генератора, когда next() называется на нем.

https://docs.python.org/3/glossary.html

При реализации интерфейса вам необходимо определить методы и сопоставить их с реализацией вашего класса. В этом случае__next__() Метод должен вызвать итератор генератора.

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

Не использовать yield в __next__ функционировать и реализовывать next также для совместимости с python2.7+

Код

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__()
Другие вопросы по тегам