Генераторы и yield в Python: большие данные без лишней памяти

Коротко о главном

Генератор — это функция с yield, которая возвращает значения по одному, по запросу. Пока вы не попросили следующий элемент, код не выполняется. Поэтому генератор не загружает всю коллекцию в память: обработка файла на гигабайты занимает примерно столько же памяти, сколько обработка одной строки.

Итератор, итерируемый объект, генератор

Список можно обходить сколько угодно раз. Итератор и генератор — только один раз: после исчерпания они пусты.

Как работает yield

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

gen = countdown(3)
next(gen)  # 3
next(gen)  # 2

Вызов countdown(3) ничего не считает — он возвращает объект-генератор. Каждый next() выполняет тело до ближайшего yield, отдаёт значение и замораживает состояние функции. Когда тело заканчивается, возбуждается StopIteration, и цикл for завершается сам.

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

Квадратные скобки создают список, круглые — генератор:

squares_list = [x * x for x in range(10_000_000)]  # всё в памяти
squares_gen = (x * x for x in range(10_000_000))   # по одному

total = sum(x * x for x in range(10_000_000))     # скобки можно опустить

Используйте генераторное выражение, когда результат нужен один раз — для sum, max, any, записи в файл.

Ленивый конвейер для большого файла

Главная сила генераторов — их можно соединять в цепочку. Каждый шаг берёт поток на вход и отдаёт поток на выход:

import csv

def read_rows(path):
    with open(path, newline="", encoding="utf-8") as f:
        yield from csv.DictReader(f)

def only_paid(rows):
    for row in rows:
        if row["status"] == "paid":
            yield row

def amounts(rows):
    for row in rows:
        yield float(row["amount"])

total = sum(amounts(only_paid(read_rows("orders.csv"))))

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

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

Для денежных сумм в реальном коде лучше Decimal, а не float, — здесь float только для краткости.

yield from

yield from iterable отдаёт все элементы другого итератора. Это короче цикла и удобно для рекурсии и разбиения генератора на части:

def walk(tree):
    yield tree.value
    for child in tree.children:
        yield from walk(child)

Кроме того, yield from корректно пробрасывает send(), исключения и возвращаемое значение вложенного генератора.

Типичные ловушки

• Одноразовость. Второй проход по тому же генератору ничего не даст:

gen = (x for x in range(3))
list(gen)  # [0, 1, 2]
list(gen)  # []

Если данные нужны дважды — создайте генератор заново или сохраните в список, если он помещается в память.

• Закрытый файл. Если вернуть генератор из функции, где файл открыт в with без yield внутри, файл закроется раньше чтения. Держите yield внутри блока with.
• len() и индексы не работают. У генератора нет длины и нельзя взять gen[5]. Для срезов есть itertools.islice.
• Ленивые ошибки. Исключение возникнет не при создании генератора, а при чтении элемента, поэтому стектрейс может указывать на неожиданное место.
• Отладка через print(gen) покажет только <generator object ...>, а не данные.

Когда генератор, а когда список

• Данные большие или бесконечные — генератор.
• Нужен один проход и агрегат — генератор.
• Нужны повторные проходы, сортировка, len, индексы — список.
• Данных мало — разница несущественна, выбирайте то, что читается проще.