Генераторы

Выражение yield

Если тело функции содержит хотя бы одно ключевое слово yield, то это так называемая функция-генератор. Вызов функции генератора не приводит к исполнению тела функции. Вместо этого функция-генератор возвращает объект генератора.

def generator_function():
    print("Начало функции.")
    yield 0

generator_object = generator_function()
type(generator_object)
generator

Объект генератора хранит в себе исходный код функции генератора, её пространство локальных переменных, а так же текущую точку выполнения. В этом и заключается главная разница между функцией и генератором.

Функция не хранит своё состояние между своими вызовами, ключевое слово return возвращает управление вызывающему коду. Ключевое слово (выражение) yield в свою очередь, передаёт управление вызывающему коду, но запоминает место и состояние (значение локальных переменных и т.п.), а не сбрасывает его. Это позволяет в последствии возобновить исполнение кода в генераторе.

Применение функции next к объекту генератора возобновляет (или начинает) исполнение кода в генерации до тех пор, пока не встретится выражение yield. Значение справа от yield возвращается функцией next вызывающему коду. Исполнение кода в генераторе ставится на паузу до следующего вызова функции next.

Note

Вообще говоря, у объекта генератора есть метод send, который позволяет не только получать объекты из генератора, но и передавать объекты из вызывающего кода в код генератора. Вызов этого метода тоже возобновляет исполнение генератора. Если g — объект генератора, то выражение g.send(None) в точности эквивалентно выражению next(g).

Продемонстрируем особенность объекта генератора на примере. Для этого напишем функцию-генератор обратного отсчета с 2.

def countdown_from_2():
    print(f"Начинаю обратный отсчет с двух.")
    yield 2
    print("Продолжаю обратный отсчет. Следующее значение 1.")
    yield 1
    print("Обратный отсчет закончен.")
    return 

g = countdown_from_2()

Выше объявлена функция-генератор countdown_from_2. И далее результата вызова этой функции связывается с именем g. Заметьте, что на экране не появилось сообщений, потому что на этот момент исполнение инструкций в тела функций ещё не началось: был создан только объект генератора.

Применим функцию next к объекту генератора.

print(f"Функция next вернула значение {next(g)}.")
Начинаю обратный отсчет с двух.
Функция next вернула значение 2.

Теперь мы видим, что инструкции в теле цикла начали исполняться, но не до конца, а только до инструкции yield 2.

def countdown_from_2():
    print(f"Начинаю обратный отсчет с двух.")
    yield 2 # <--- pause
    print("Продолжаю обратный отсчет. Следующее значение 1.")
    yield 1
    print("Обратный отсчет закончен.")
    return 

Применим функцию next ещё раз.

print(f"Функция next вернула значение {next(g)}.")
Продолжаю обратный отсчет. Следующее значение 1.
Функция next вернула значение 1.

Инструкции продолжились исполняться, пока не встретился очередной yield.

def countdown_from_2():
    print(f"Начинаю обратный отсчет с двух.")
    yield 2 
    print("Продолжаю обратный отсчет. Следующее значение 1.")
    yield 1 # <--- pause
    print("Обратный отсчет закончен.")
    return 

Последний раз применим функцию next. Заметим, что дальше по пути исполнения программы встречается ключевое слово return, а не yield.

def countdown_from_2():
    print(f"Начинаю обратный отсчет с двух.")
    yield 2 
    print("Продолжаю обратный отсчет. Следующее значение 1.")
    yield 1 
    print("Обратный отсчет закончен.")
    return # <--- raise StopIteration

Ключевое слово return в генераторе возбуждает исключение StopIteration, что приводит к выходу из тела функции генератора.

try:
    next(g)
except StopIteration:
    print("Генератор исчерпан.")
Обратный отсчет закончен.
Генератор исчерпан.

Итераторы vs генераторы

Функции генераторы удобно задействовать для создания итераторов, т.к. они поддерживают тот же интерфейс: по запросу функции next выдаётся очередное значение, по исчерпании элементов возбуждается исключение StopIteration.

В качестве примера рассмотрим наивную реализацию альтернативы range, но для действительных чисел.

def frange(start, stop, step=1.0):
    while start < stop:
        yield start
        start += step

for x in frange(0, 1, 0.1):
    print(x, end=" ")
0 0.1 0.2 0.30000000000000004 0.4 0.5 0.6 0.7 0.7999999999999999 0.8999999999999999 0.9999999999999999 

Ключевое отличие генератора от классического итератора заключается в том, что итератор выдаёт уже существующие в каком-то контейнере значения, а генератор вычисляет новые значения на лету. Это позволяет экономить ресурсы системы, если для дальнейших вычислений не требуются, чтобы все значения где-то хранились в одном месте.

В качестве примера рассмотрим вычисление числа \(\pi\) через сумму ряда

\[ \sum_{n=1}^\infty \dfrac{1}{n^2} = \dfrac{\pi^2}{6}. \]
from math import sqrt

def list_of_terms(N):
    return [1./(n * n) for n in range(1, N)]


def generator_of_terms(N):
    for n in range(1, N):
        yield 1/(n * n)


def pi_from_sum(S):
    return sqrt(6*S)


N = 1_000_000
S1 = sum(list_of_terms(N))
S2 = sum(generator_of_terms(N))

print(pi_from_sum(S1), pi_from_sum(S2))
3.141591698659554 3.141591698659554

Профилирование по памяти продемонстрирует, что функция с генератором гораздо экономнее при больших N, чем функция со списком. Это объясняется тем, что ни в один момент времени не создаётся список, чтобы хранить члены ряда. Вместо этого, они вычисляются по запросу функции sum. Такой подход, когда вычисления откладываются до тех пор, пока не потребуется их результат, называют ленивыми вычислениями.

Т.к. генератор поддерживает протокол итерации, то при необходимости можно получить список из генератора.

print(list(frange(0, 1, 0.1)))
[0, 0.1, 0.2, 0.30000000000000004, 0.4, 0.5, 0.6, 0.7, 0.7999999999999999, 0.8999999999999999, 0.9999999999999999]

Но это может привести к зацикливанию, если генератор никогда не исчерпается. В этом заключается ещё одно крупное отличие генераторов от итераторов: генераторы могут генерировать бесконечную последовательность элементов, а итераторы всегда пробегаются по расположенной в памяти, а значит конечной, последовательности элементов. Например, следующий генератор выдаёт бесконечную последовательность натуральных чисел.

Note

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

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

Т.к. такой генератор никогда не исчерпается, то он никогда и не бросит исключения StopIteration, а значит применять его в цикле for можно только в случае, если предусмотрен выход оператором break. Иначе программа зациклится.

Генераторные выражения

Если вместо квадратных скобочек в списковом включении указать круглые, то вы получите генераторное выражение. Разница заключается в том, что в случае спискового включения все вычисления производятся сразу и в результате выходит список, а в случае генераторного выражения вы получаете объект генератора, что само по себе не приводит к никаким вычислениям.

Чтобы продемонстрировать это, создадим список и генератор схожим выражением и измерим, сколько байт занимает каждый из них.

Note

Метод getsizeof измеряет количество байт, которое занимает объект в памяти. Он корректно работает для всех встроенных объектов и объектов из стандартной библиотеки, но может давать ложную информацию для пользовательских объектов и объектов из сторонних библиотек.

from sys import getsizeof

l = [x * x for x in range(1_000_000)]
g = (x * x for x in range(1_000_000))

print(f"Список занимает {getsizeof(l)} байт")
print(f"Генератор занимает {getsizeof(g)} байт")
Список занимает 8697456 байт
Генератор занимает 112 байт