Генераторы Python: функции с yield и ленивые вычисления

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

Генераторная функция и yield

Функция становится генераторной, если в ее теле есть ключевое слово yield. При вызове такая функция не выполняет свой код, а возвращает объект-генератор. Код начинает работать только тогда, когда у генератора запрашивают значение функцией next() или циклом for.

def countdown(n):
    while n > 0:
        yield n
        n -= 1

gen = countdown(3)
print(gen)        # <generator object countdown at 0x...>
print(next(gen))  # 3
print(next(gen))  # 2
print(next(gen))  # 1
next(gen)         # StopIteration

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

for value in countdown(3):
    print(value, end=" ")  # 3 2 1

Как выполняется генератор

Порядок выполнения хорошо видно, если добавить вывод в тело функции. Каждый next() прогоняет код только до ближайшего yield.

def steps():
    print("старт")
    yield 1
    print("после первого yield")
    yield 2
    print("конец функции")

g = steps()     # ничего не напечатано
next(g)         # старт
next(g)         # после первого yield
next(g, None)   # конец функции, вернет None вместо StopIteration

Из этого следуют три правила, о которых важно помнить:

• Генератор одноразовый. После того как значения закончились, повторный цикл по тому же объекту ничего не выдаст. Чтобы пройти заново, вызовите функцию еще раз.
• У генератора нет длины и индексов: len(gen) и gen[0] вызовут TypeError. Если нужен список, используйте list(gen).
• Инструкция return в генераторе завершает его. Возвращенное значение попадает в атрибут value исключения StopIteration, в цикле for оно не видно.

Выражения-генераторы

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

import sys

squares_list = [x * x for x in range(1_000_000)]
squares_gen = (x * x for x in range(1_000_000))

print(sys.getsizeof(squares_list))  # около 8 МБ
print(sys.getsizeof(squares_gen))   # около 100-200 байт
print(sum(x * x for x in range(1_000_000)))  # 333332833333500000

Список сразу вычисляет и хранит миллион чисел. Генератор занимает в памяти одинаково мало при любом размере, потому что хранит только текущее состояние. Если выражение-генератор единственный аргумент функции, вторые скобки можно не писать, как в вызове sum() выше.

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

Где генераторы полезны

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

def read_errors(path):
    with open(path, encoding="utf-8") as f:
        for line in f:
            if "ERROR" in line:
                yield line.rstrip()

for line in read_errors("app.log"):
    print(line)

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

lines = read_errors("app.log")
codes = (line.split()[-1] for line in lines)
unique = set(codes)

Генератор может быть и бесконечным: цикл while True с yield внутри никогда не завершится сам, значения берут ровно столько, сколько нужно. Ограничить их помогает itertools.islice, о нем и других инструментах в статье про модуль itertools.

from itertools import islice

def fibonacci():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

print(list(islice(fibonacci(), 10)))
# [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

yield from и методы генератора

Конструкция yield from передает наружу все значения другого итерируемого объекта. Она заменяет вложенный цикл for x in other: yield x и особенно удобна в рекурсии.

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

print(list(flatten([1, [2, [3, 4]], 5])))  # [1, 2, 3, 4, 5]

У объекта-генератора есть методы для управления им извне. send(value) продолжает выполнение и передает значение, которое становится результатом выражения yield внутри функции. close() останавливает генератор досрочно: внутри возникает GeneratorExit, и блоки finally выполняются. throw() выбрасывает исключение в точке остановки. В обычном коде эти методы нужны редко, на них построены старые сопрограммы и некоторые библиотеки.

def running_average():
    total = count = 0
    average = None
    while True:
        value = yield average
        total += value
        count += 1
        average = total / count

avg = running_average()
next(avg)            # запуск до первого yield
print(avg.send(10))  # 10.0
print(avg.send(20))  # 15.0

Генераторы и итераторы

Любой генератор является итератором: у него есть методы __iter__() и __next__(). Разница в способе создания. Итератор-класс требует описать оба метода и хранить состояние в атрибутах, а генератор дает тот же результат обычной функцией, состояние сохраняется автоматически.

class Countdown:
    def __init__(self, n):
        self.n = n
    def __iter__(self):
        return self
    def __next__(self):
        if self.n <= 0:
            raise StopIteration
        self.n -= 1
        return self.n + 1

print(list(Countdown(3)))   # [3, 2, 1]
print(list(countdown(3)))   # [3, 2, 1], тот же результат функцией из начала статьи

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

Частые вопросы

Короткие ответы на вопросы, которые возникают при изучении генераторов.

Чем yield отличается от return?

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

Как проверить, что объект является генератором?

Используйте inspect.isgenerator(obj) для объекта и inspect.isgeneratorfunction(func) для функции. Также подойдет isinstance(obj, types.GeneratorType).

Почему второй цикл по генератору ничего не выводит?

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

Быстрее ли генератор, чем список?

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