Pixamp

Генераторы в Python

18.03.2019 22:00 CPython 3.7.2

The article image

Генераторы, как и итераторы, являются неотъемлемой частью Python. Понимание и умение их использовать в работе важно для каждого Python-программиста. Итераторы уже были главной темой статьи Итераторы и итерируемые объекты в Python, а сегодня пришло время поговорить о генераторах.

Данная статья будет полезна как новичкам, так и более опытным программистам, желающим расширить свои знания по данному вопросу.

Примечание

Если вы еще не прочитали статью Итераторы и итерируемые объекты в Python, то я рекомендую сделать это перед чтением данной статьи.

Генераторы

Начнем с формального определения: генератор (generator) – это функция, которая возвращает специальный итератор. Такую функцию еще называют функцией генератора (generator function), а возвращаемый итератор – итератором генератора (generator iterator). В отличие от обычной функции Python, в теле функции генератора можно встретить минимум одну конструкцию yield.

Примечание

Под термином “генератор” зачастую понимается именно функция генератора, но иногда подразумевается то, что возвращает эта функция, а именно – итератор генератора.

По традиции начнем с примера. Предположим, что мы хотели бы иметь под рукой такой итератор, который возвращал бы произвольное количество степеней двойки (1, 2, 4, 8, 16 и так далее).

Напишем класс итератора. Напомню, что мы должны реализовать методы __iter__() и __next__():

class TwoPowerful:
    def __init__(self, number):
        self.number = number

    def __iter__(self):
        self.range = iter(range(self.number))
        return self

    def __next__(self):
        return 2 ** next(self.range)

Попробуем его в деле:

>>> for power in TwoPowerful(5):
...     print(power)
...
1
2
4
8
16

Все работает! Нужное количество степеней задается положительным целым числом при создании экземпляра класса TwoPowerful. Обратите внимание, что выше я использовал класс range для того, чтобы получить последовательность показателей степени (в данном примере – это последовательность от “0” до “4” включительно). Так как объект range является итерируемым объектом, а не итератором – я использую встроенную функцию iter() для того, чтобы получить итератор.

Эту же задачу мы можем решить, используя генератор. Как я уже писал выше, генератор – это функция, возвращающая итератор. В теле такой функции можно найти минимум одну конструкцию yield:

>>> def two_powerful(number):
...     for power in range(number):
...         yield 2 ** power
...
>>> for power in two_powerful(5):
...     print(power)
...
1
2
4
8
16

Мы получили тот же результат, но использовали другой подход. Первое, что бросается в глаза – код стал компактнее: одна небольшая функция вместо сравнительно громоздкого класса итератора с тремя методами. Также код стал более читабельным, не так ли?


Используя генераторы, мы можем постепенно обрабатывать большие последовательности данных, объем которых может значительно превышать доступный объем памяти. Это могут быть, например, большие файлы или базы данных. Также мы можем генерировать бесконечные последовательности данных (например, последовательность простых чисел или чисел Фибоначчи).

Наша последовательность степеней двойки тоже может быть бесконечной, достаточно слегка модифицировать функцию генератора:

from itertools import count

def two_powerful(number=0):
    powers = range(number) if number else count()
    for power in powers:
        yield 2 ** power

Если number имеет значение 0 (на самом деле вместо 0 может быть None, пустая строка, пустой список и тому подобное), то вместо класса range используется функция count() из модуля itertools. В данном случае count() генерирует бесконечную последовательность целых чисел, начиная с нуля (0, 1, 2, 3, 4…), а это как раз то, что нам нужно.

Проверим обновленный генератор:

>>> for i, power in enumerate(two_powerful()):
...     print(power)
...     if i == 4:  # выйти из цикла после 5-го по счету значения
...         break
...
1
2
4
8
16

Конструкция yield

Давайте разберемся, в чем же отличие функции генератора, от обычной функции.

Рассмотрим две функции:

def nongenerator(x):
    return x

def generator(x):
    yield x

Здесь nongenerator() является обычной функцией, а generator() – функцией генератора.

Давайте, для начала, посмотрим на типы наших объектов (функции в Python тоже являются объектами):

>>> type(nongenerator)
<class 'function'>
>>> type(generator)
<class 'function'>
>>> type(nongenerator) is type(generator)
True

Как видно, здесь разницы никакой нет. В обоих случаях мы видим один и тот же тип – function. Но давайте теперь вызовем эти функции и посмотрим на типы результатов:

>>> arg = 5
>>> type(arg)
<class 'int'>
>>> type(nongenerator(arg))
<class 'int'>
>>> type(generator(arg))
<class 'generator'>
>>> type(nongenerator(arg)) is type(generator(arg))
False

Я вызываю обе функции с аргументом arg, при этом nongenerator(), ожидаемо, возвращает объект типа int, а вот generator() – объект типа generator.

Посмотрим еще раз на обе функции: разница между ними только в том, что nongenerator() в теле содержит return, а generator()yield. Именно наличие yield и приводит к тому, что функция начинает возвращать объект типа generator. Про этот объект я уже писал выше - это и есть итератор генератора. Этот итератор мы можем передать функции next() или, например, использовать в цикле for:

>>> arg = 5
>>> next(generator(arg))
5
>>> for i in generator(arg):
...     print(i)
...
5
>>> type(nongenerator(arg)) is type(next(generator(arg)))
True

Вернемся теперь к примеру со степенями двойки:

def two_powerful(number):
    for power in range(number):
        yield 2 ** power

Здесь ситуация несколько сложнее, так как yield находится в теле цикла for. Напомню результат работы данной функции с аргументом 5:

>>> for power in two_powerful(5):
...     print(power)
...
1
2
4
8
16

Какой результат мы получим, если сейчас заменить yield на return? Давайте проверим:

>>> def two_powerless(number):
...     for power in range(number):
...         return 2 ** power
...
>>> for power in two_powerless(5):
...     print(power)
...
Traceback (most recent call last):
    ...
TypeError: 'int' object is not iterable

Мы получили ошибку! Все потому, что two_powerless(), в отличие от two_powerful(), возвращает целое число, а не итератор:

>>> two_powerless(5)
1

Функция two_powerless() возвращает только 1 – первую степень двойки.

Отсюда мы можем сделать следующие выводы:

  1. Как только будет достигнут return, выполнение кода обычной функции прекратится, а результат работы будет возвращен в точку вызова.
  2. В случае же yield работа функции не прекращается, а временно прерывается, возвращается текущее значение выражения yield (yield <выражение>), а затем выполнение кода функции продолжается. Значений будет возвращено ровно столько, сколько будет встречена конструкция yield:

Yield expression

two_powerful() возвращает объект генератора, который является своего рода оберткой над этой функцией. Когда мы вызываем next() первый раз, передавая ей полученный объект генератора, код two_powerful() начинает выполняться до тех пор, пока не будет встречена конструкция yield, после чего выполнение будет приостановлено, а next() вернет значение выражения yield в точку вызова. Следующий вызов next() приведет к тому, что выполнение функции продолжится с места остановки, пока снова не будет встречена yield (yield находится в теле цикла) и так далее:

>>> g = two_powerful(5)
>>> next(g)
1
>>> next(g)
2
>>> next(g)
4
>>> next(g)
8
>>> next(g)
16

В конечном счете, range будет истощен и for прекратит работу.

Как мы помним, что как только итератор истощен, должно быть сгенерировано исключение StopIteration. Использование return в теле функции генератора и приводит к генерации StopIteration. В примере выше return явно отсутствует в теле функции two_powerful(), тем не менее в Python любая функция всегда возвращает какое-нибудь значение. В нашем случае – это значение None: интерпретатор добавляет недостающий return в конец функции. Выглядит это примерно так:

def two_powerful(number):
    for power in range(number):
        yield 2 ** power
    return None

Шестой вызов next() приведет к генерации StopIteration:

>>> next(g)
Traceback (most recent call last):
    ...
StopIteration

Обратите внимание, что возвращаемое значение используется при генерации StopIteration:

>>> def two_powerful(number):
...     for power in range(number):
...         yield 2 ** power
...     return 'I\'m so exhausted!'
...
>>> g = two_powerful(5)
>>> next(g)
1
>>> next(g)
2
>>> next(g)
4
>>> next(g)
8
>>> next(g)
16
>>> next(g)
Traceback (most recent call last):
    ...
StopIteration: I'm so exhausted!

Проверим атрибут value объекта исключения (в примере рассматривается последний next()):

>>> try:
...     next(g)
... except StopIteration as e:
...     e.value
...
"I'm so exhausted!"

Фактически, это еще один способ возвращать данные. Используя блок try/except, мы можем получить доступ к возвращаемому значению функции.


В примерах выше конструкция yield всегда была в “хвосте” функции, но, конечно же, месторасположение может быть любым:

>>> def generator(x):
...     print('Hello')
...     yield x
...     print('Goodbye')
...
>>> for i in generator(5):
...     print(i)
...
Hello
5
Goodbye

При необходимости в теле функции может быть больше одной конструкции yield:

>>> def generator(x):
...     yield x + 1
...     yield x + 2
...     yield x + 3
...
>>> for i in generator(5):
...     print(i)
...
6
7
8

Стоит также заметить, что функция будет возвращать итератор генератора независимо от того достижима конструкция yield или нет:

>>> def generator(x):
...     print('Hello')
...     if False:
...         yield x
...     print('Goodbye')
...
>>> g = generator(5)
>>> type(g)
<class 'generator'>
>>> next(g)
Hello
Goodbye
Traceback (most recent call last):
    ...
StopIteration

Несмотря ни на что, функция generator() по-прежнему возвращает итератор, который генерирует StopIteration при первом же next().

Метод send()

Давайте рассмотрим несколько весьма полезных методов объекта генератора. Один из них – метод send(). Он позволяет передавать значения функции генератора.

Напомню, что в предыдущих примерах мы только получали значения выражения yield. Графически это можно представить так:

One-way generator

Вот как меняется ситуация, когда мы начинаем отправлять значения, используя send():

Two-way generator

Давайте рассмотрим простой пример:

def generator(x):
    while True:
        x = yield x

Из нового для нас здесь то, что теперь yield находится справа от =.

Давайте проверим эту функцию в деле:

>>> g = generator(5)
>>> g.send(None)
5
>>> g.send(10)
10
>>> g.send(15)
15
>>> g.send(20)
20

Обратите внимание, что здесь я не использую функцию next(), так как это попросту не требуется – send() сам выполняет всю нужную работу. Вызов send(None) аналогичен вызову next(), то есть мы могли бы написать так:

>>> g = generator(5)
>>> next(g)  # аналогично g.send(None)
5
>>> g.send(10)
10
>>> g.send(15)
15
>>> g.send(20)
20

Что же происходит, когда мы вызываем метод send() итератора генератора? Первый вызов (send(None)), как и в случае с next(), приводит к тому, что код функции generator() начинает выполняться до тех пор, пока не будет встречено выражение yield, после чего выполнение данной функции приостанавливается, а значение выражения yield (а это – 5) возвращается в точку вызова send(). На этом этапе переменной x еще не присвоено никакое значение. Благодаря следующему вызову send() (send(10)) выполнение функции generator() продолжится с места остановки и вот теперь x будет присвоено 10 (представьте, что x = yield x как бы меняется на x = 10). Так как yield находится в теле бесконечного цикла, то мы можем бесконечно вызывать send(), передавая ей разные аргументы, которые тут же будут возвращаться назад.

Почему мы передаем None (send(None)) в самом начале? Все потому, что это первый запуск и выражение yield еще не было встречено.

Примечание

Этот первый вызов (send(None) или next()) в зарубежной литературе часто называют “priming”.

Кстати, имейте ввиду, что нам не разрешат передать что-то отличное от None:

>>> g = generator(5)
>>> g.send(10)  # вместо send(None)
Traceback (most recent call last):
    ...
TypeError: can't send non-None value to a just-started generator

Но давайте рассмотрим еще парочку похожих примеров. Изменим слегка нашу функцию:

def generator(x):
    while True:
        x = yield x + 1

Мы видим, что выражение yield немного изменилось: yield x + 1 вместо yield x. Какой результат мы получим в этом случае? Давайте посмотрим:

>>> g = generator(5)
>>> g.send(None)
6
>>> g.send(10)
11
>>> g.send(15)
16
>>> g.send(20)
21

Вполне ожидаемо, не так ли? Но что произойдет, если изменить функцию вот таким образом:

def generator(x):
    while True:
        x = (yield x) + 1

Посмотрим на результат:

>>> g = generator(5)
>>> g.send(None)
5
>>> g.send(10)
11
>>> g.send(15)
16
>>> g.send(20)
21

Как видно, в этот раз вместо 6, при первом вызове send(), мы получили 5: единица прибавляется только после возобновления работы функции. Здесь, как мне кажется, хорошо виден механизм приостановки/возобновления выполнения функции генератора.


Если требуется только передача значений, то можно использовать “голое” выражение yield (в этом случае, как и с “голым” return, будет возвращаться None):

import sys

def echo(end='\n', file=sys.stdout):
    while True:
        s = yield
        file.write(f'{s}{end}')

Функция echo(), подобно print(), пишет данные в файл (по умолчанию это stdout):

>>> e = echo()
>>> e.send(None)
>>> e.send('Hello')
Hello

Как вы уже наверняка знаете, обычные функции Python называют подпрограммами (subroutines). Но функция echo(), из примера выше, называется сопрограммой (coroutine). Сопрограмма - это такая функция, которая может производить множество значений, работа ее может приостанавливаться и затем снова возобновляться, а также ей можно передавать значения.

Примечание

Начиная с версии 2.5 в Python появилась возможность писать такие сопрограммы. Впрочем, только в версии 3.5 был добавлен специальный синтаксис для создания сопрограмм (async/await), но об этом мы поговорим в следующих статьях.

Для того, чтобы не писать каждый раз e.send(None) можно реализовать специальный декоратор (про декораторы мы поговорим в следующих статьях), который будет выполнять подготовительные работы:

from functools import wraps

def coroutine(func):
    @wraps(func)
    def prime(*args,**kwargs):
        c = func(*args,**kwargs)
        c.send(None)
        return c
    return prime

И тогда предыдущий пример с сопрограммой echo() можно переписать вот так:

@coroutine
def echo(end='\n', file=sys.stdout):
    while True:
        s = yield
        file.write(f'{s}{end}')

Попробуем функцию в деле:

>>> e = echo()
>>> e.send('Hello')  # send(None) больше ненужен
Hello

Метод close()

Рассмотрим пример:

def count(start=0, step=1):
    n = start
    while True:
        yield n
        n += step

Функция count() позволяет нам генерировать бесконечные числовые последовательности, начиная со start и с шагом step (по умолчанию последовательность начинается с нуля, а шаг равен единице):

>>> c = count()
>>> next(c)
0
>>> next(c)
1
>>> next(c)
2
>>> next(c)
3

В теле функции count() мы видим while True, поэтому генерируемая последовательность будет бесконечной: генератор будет возвращать все новые и новые значения при вызове next(). Это может быть полезно в ряде случаев, но как сказать ему “Горшочек, больше не вари!”, если мы хотим прекратить генерацию новых значений? Для этого существует метод close():

>>> c = count()
>>> next(c)
0
>>> c.close()
>>> next(c)  # на этом этапе итератор генератора уже "истощен"
Traceback (most recent call last):
    ...
StopIteration

Когда мы вызываем close(), генерируется исключение GeneratorExit в точке остановки функции. Перехватывать это исключение не нужно, никакого сообщения об ошибке мы не получим: close() все сделает сам.

Давайте немного модифицируем нашу функцию:

def count(start=0, step=1):
    print('Hello')
    n = start
    while True:
        yield n
        n += step
    print('Goodbye')

Я добавил два print() в начало и в конец тела функции, чтобы было проще понять что происходит, когда мы вызываем close():

>>> c = count()
>>> next(c)
Hello
0
>>> c.close()
>>> next(c)
Traceback (most recent call last):
    ...
StopIteration

Сообщение Goodbye мы так и не увидели, что логично – функция прекратила свою работу в точке остановки (yield n). Зная это, мы можем модифицировать наш код таким образом:

def count(start=0, step=1):
    print('Hello')
    n = start
    try:
        while True:
            yield n
            n += step
    except GeneratorExit:
        print('Goodbye')

Теперь исключение GeneratorExit перехватывается в блоке try/except. Какой результат мы получим в этом случае? Давайте посмотрим:

>>> c = count()
>>> next(c)
Hello
0
>>> c.close()
Goodbye
>>> next(c)
Traceback (most recent call last):
    ...
StopIteration

Мы перехватили GeneratorExit и наконец-то увидели Goodbye!

Как вы думаете, что произойдет, если добавить еще одну конструкцию yield ниже блока try/except?

def count(start=0, step=1):
    print('Hello')
    n = start
    try:
        while True:
            yield n
            n += step
    except GeneratorExit:
        print('Goodbye')
    yield

Как мы увидим дальше – это недопустимо. В ответ мы получим RuntimeError:

>>> c = count()
>>> next(c)
Hello
0
>>> c.close()
Goodbye
Traceback (most recent call last):
    ...
RuntimeError: generator ignored GeneratorExit

Обратите внимание, что второй и последующий вызовы close() по сути уже ничего не делают:

>>> c = count()
>>> next(c)
Hello
0
>>> c.close()
>>> c.close()
>>> c.close()

Метод throw()

Любое исключение сгенерированное во время работы функции генератора (которое не было перехвачено), приведет к истощению итератора, а само исключение будет поднято наверх.

Рассмотрим пример:

def count(start=0, step=1):
    n = start
    while True:
        yield n
        n += step
        raise GeneratorExit

Здесь мы явно генерируем GeneratorExit. Давайте посмотрим, какой будет результат:

>>> c = count()
>>> next(c)
0
>>> next(c)
Traceback (most recent call last):
  File "<input>", line 1, in <module>
    next(c)
  File "<input>", line 6, in count
    raise GeneratorExit
GeneratorExit
>>> next(c)  # итератор генератора уже истощен
Traceback (most recent call last):
    ...
StopIteration

Так как сейчас мы не используем метод close(), то GeneratorExit поднимается наверх. Что касается итератора, то он перестает возвращать новые значения.

Используя метод throw(), можно генерировать исключения в точке остановки функции. Чтобы продемонстрировать работу этого метода, давайте немного изменим нашу функцию:

def count(start=0, step=1):
    print('Hello')
    n = start
    while True:
        yield n
        n += step
    print('Goodbye')

Попробуем сгенерировать то же самое исключение GeneratorExit, но теперь с помощью throw():

>>> c = count()
>>> next(c)
Hello
0
>>> c.throw(GeneratorExit)
Traceback (most recent call last):
  File "<input>", line 1, in <module>
    c.throw(GeneratorExit)
  File "<input>", line 5, in count
    yield n
GeneratorExit
>>> next(c)
Traceback (most recent call last):
    ...
StopIteration

Мы не увидели Goodbye, так как print('Goodbye') располагается ниже yield n (а именно это и есть “точка остановки”). Впрочем, ничего не мешает нам перехватить это исключение в теле функции:

def count(start=0, step=1):
    print('Hello')
    n = start
    try:
        while True:
            yield n
            n += step
    except GeneratorExit:
        print('Goodbye')

Посмотрим на результат:

>>> c = count()
>>> next(c)
Hello
0
>>> c.throw(GeneratorExit)
Goodbye
Traceback (most recent call last):
    ...
StopIteration

Теперь при вызове throw() мы видим как Hello, так и Goodbye, но в то же время получаем StopIteration. Это все потому, что throw() после генерации исключения еще пытается получить следующее значение (как мы это делали с помощью next()), но возвращать уже нечего. Чтобы убедиться в этом, изменим еще раз нашу функцию:

def count(start=0, step=1):
    print('Hello')
    n = start
    try:
        while True:
            yield n
            n += step
    except CustomException:
        print('Goodbye')
    yield

Я добавил yield в конец функции, а также, для разнообразия, заменил GeneratorExit на CustomException:

class CustomException(Exception):
    pass

Вот какой результат мы получим в этот раз:

>>> c = count()
>>> next(c)
Hello
0
>>> c.throw(CustomException)
Goodbye
>>> next(c)
Traceback (most recent call last):
    ...
StopIteration

throw() в этот раз возвращает None благодаря последней конструкции yield. А вот следующий вызов next() вполне ожидаемо приводит к генерации StopIteration.

Вообще, чтобы гарантированно выполнить нужный завершающий код лучше всего использовать try/finaly:

def count(start=0, step=1):
    print('Hello')
    try:
        n = start
        while True:
            yield n
            n += step
    finally:
        print('Goodbye')

В этом случае мы увидим Goodbye при вызове throw():

>>> c = count()
>>> next(c)
Hello
0
>>> c.throw(CustomException)
Goodbye
Traceback (most recent call last):
    ...
CustomException

A также при вызове close():

>>> c = count()
>>> next(c)
Hello
0
>>> c.close()
Goodbye

Но что произойдет, если вызвать throw() сразу после получения итератора генератора? Давайте проверим:

>>> c = count()
>>> c.throw(CustomException)
Traceback (most recent call last):
  File "<input>", line 1, in <module>
    c.throw(CustomException)
  File "<input>", line 1, in count
    def count(start=0, step=1):
CustomException
>>> next(c)
Traceback (most recent call last):
    ...
StopIteration

CustomException по-прежнему генерируется, но так как yield еще не было достигнуто на этом этапе, то в Traceback мы видим другую точку генерации исключения.

Конструкция yield from

Начиная с версии 3.3, в Python был добавлен новый синтаксис для yieldyield from. С этого момента отпала необходимость использовать, например, цикл for для того, чтобы получить значения одного генератора в теле функции другого.

Рассмотрим пример простого генератора, который возвращает числа от нуля до x:

def from_zero_to_any(x):
    for i in range(x):
        yield i

При x = 5 мы получим такой результат:

>>> for i in from_zero_to_any(5):
...     print(i)
...
0
1
2
3
4

Используя конструкцию yield from, мы можем сократить код функции from_zero_to_any():

def from_zero_to_any(x):
    yield from range(x)  # цикл for больше ненужен

При этом конечный результат не изменится:

>>> for i in from_zero_to_any(5):
...     print(i)
...
0
1
2
3
4

Чтобы понять, что происходит при использовании yield from, давайте для начала вспомним, что происходит, когда мы используем просто yield (не yield from):

Two-way generator

Здесь мы видим, что итератор генератора (на схеме обозначен как Generator) возвращает значения в точку вызова (Caller). В свою очередь Caller передает значения Generator, используя метод send(). Также Caller может остановить генерацию новых значений (“закрыть генератор”) с помощью метода close(), либо сгенерировать исключение в точке остановки функции, используя метод throw().

А вот что происходит при использовании yield from:

Yield from

Модуль inspect

Напоследок хотелось бы упомянуть несколько полезных функций модуля inspect из стандартной библиотеки Python. Используя данные функции, мы можем посмотреть текущее состояние итератора генератора, проверить является ли функция функцией генератора, а также является ли объект итератором генератора.

Состояние итератора генератора можно проверить с помощью функции getgeneratorstate():

>>> import inspect
>>> def generator(x):
...     while x != 0:
...         x = yield x
...
>>> g = generator(5)
>>> inspect.getgeneratorstate(g)
'GEN_CREATED'
>>> g.send(None)
5
>>> inspect.getgeneratorstate(g)
'GEN_SUSPENDED'
>>> g.send(10)
10
>>> inspect.getgeneratorstate(g)
'GEN_SUSPENDED'
>>> g.send(0)
Traceback (most recent call last):
    ...
StopIteration
>>> inspect.getgeneratorstate(g)
'GEN_CLOSED'

Функция getgeneratorstate() может вернуть четыре возможных состояния генератора:

  1. GEN_CREATED – состояние ожидания выполнения.
  2. GEN_RUNNING – выполняется в данный момент.
  3. GEN_SUSPENDED – выполнение приостановлено.
  4. GEN_CLOSED – выполнение завершено.

Для того, чтобы проверить, является ли некая функция функцией генератора используется isgeneratorfunction():

>>> inspect.isgeneratorfunction(generator)
True

А для проверки является ли какой-либо объект итератором генератора – isgenerator():

>>> inspect.isgenerator(g)
True

Список полезных ссылок по теме:

  1. The yield statement
  2. Yield expressions
  3. PEP 380 – Syntax for Delegating to a Subgenerator

Дата последнего редактирования статьи: 28.04.2020 14:25

Предыдущая статья › Итераторы и итерируемые объекты в Python

Главная страница