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

Категория

Ленивые вычисления

Уровень

Продвинутый

Связанные темы

Итераторы, функции, включения, yield

Зачем это знать

Обрабатывать большие и потоковые данные по частям

§1Генератор — разновидность итератора

Генератор выдаёт очередное значение только тогда, когда его запрашивают. Он поддерживает iter(), next() и заканчивается через StopIteration, поэтому полностью следует протоколу итератора.

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

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

Наличие yield меняет поведение функции: её вызов возвращает генератор, а значения появляются при переборе.

def numbers():
    yield 1
    yield 2
    yield 3

generator = numbers()
print(next(generator))
print(next(generator))
print(next(generator))

1
2
3

Решение принимается по тексту функции, а не во время работы: если Python видит в теле yield, вся функция становится генераторной — независимо от того, выполнится эта строка или нет.

§3Тело не выполняется при создании

Вызов генераторной функции лишь создаёт генератор. Первая строка тела выполняется при первом запросе значения.

def demo():
    print("начало")
    yield 10

generator = demo()
print("создан")
print(next(generator))

создан
начало
10

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

§4yield возвращает значение и ставит паузу

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

def count_up(limit):
    current = 0
    while current < limit:
        yield current
        current += 1

generator = count_up(3)
print(next(generator))
print(next(generator))
print(next(generator))

0
1
2

Из схемы видно главное: две части программы работают по очереди, а не одновременно. Пока идёт закрашенный отрезок, вызывающий код ждёт; как только генератор доходит до yield, он замирает — и в этот момент его переменная current продолжает существовать, сохраняя значение до следующего next(). Внизу схемы показано, чему она равна на каждой паузе.

Это и есть главное отличие от обычной функции. Обычная функция после return забывает всё: её локальные переменные исчезают. Приостановленный генератор продолжает существовать вместе с ними — current живёт между вызовами и помнит своё значение.

§5Обычный способ чтения — цикл for

Цикл сам запрашивает значения и останавливается в конце генератора.

def squares(limit):
    for number in range(limit):
        yield number ** 2

for value in squares(5):
    print(value)

0
1
4
9
16

Ручной next() из предыдущих параграфов нужен был, чтобы разглядеть механику. В рабочем коде генератор почти всегда читают циклом: он сам перехватывает StopIteration и не даёт ошибиться на границе.

§6Генераторное выражение

Оно похоже на списковое включение из главы 4.10, но использует круглые скобки.

squares = (number ** 2 for number in range(5))
print(list(squares))

[0, 1, 4, 9, 16]

Одна пара скобок меняет всё: [...] строит список немедленно, (...) создаёт генератор и не вычисляет ничего. Условие if в конце работает так же, как во включениях: (x for x in data if x > 0).

§7Список и генераторное выражение

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

§8Сколько это стоит в памяти

Разницу удобно увидеть в числах — функция sys.getsizeof() показывает размер объекта в байтах.

import sys

squares_list = [n * n for n in range(1000000)]
squares_gen = (n * n for n in range(1000000))

print(sys.getsizeof(squares_list))
print(sys.getsizeof(squares_gen))

8448728
200

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

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

§9Генератор исчерпывается

generator = (number for number in range(3))
print(list(generator))
print(list(generator))

[0, 1, 2]
[]

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

§10Потенциально бесконечный генератор

Генератор может описывать последовательность без заранее известного конца. Потребитель обязан сам ограничить чтение.

def natural_numbers():
    number = 1
    while True:
        yield number
        number += 1

generator = natural_numbers()
for _ in range(4):
    print(next(generator))

1
2
3
4

Списком такое не выразить в принципе: бесконечный список построить нельзя. Здесь же while True безопасен — цикл делает ровно столько витков, сколько значений запросили.

§11Передача генератора функции

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

result = sum(number ** 2 for number in range(10))
print(result)

285

Если генераторное выражение — единственный аргумент функции, его собственные круглые скобки можно не дублировать. Так же устроены max(), min(), any(), all() и join(): все они читают значения по одному и не нуждаются в готовом списке.

§12Построчная обработка файлов

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

def non_empty_lines(file):
    for line in file:
        cleaned = line.strip()
        if cleaned:
            yield cleaned

Открывать файл лучше через with, чтобы он гарантированно закрылся после обработки. Это самый частый рабочий сценарий генераторов: файл на несколько гигабайт обрабатывается программой, которая занимает несколько мегабайт памяти. Полную работу с файлами разберём в модуле 8, а особенности with внутри генератора — в главе 7.4.

§13return и yield

return без значения внутри генераторной функции завершает последовательность. Возвращать значение он тоже может, но при обычном переборе оно не видно — забрать его удаётся только делегированием из главы 7.5.

§14Когда генератор не нужен

Ленивость полезна не всегда, и выбор стоит делать осознанно.

§15Частые ошибки

§16Проверьте себя

§17Шпаргалка

Итог главы

Генератор — итератор, создаваемый функцией с yield или генераторным выражением. Он начинает вычисления при запросе, сохраняет состояние между значениями и исчерпывается после одного полного прохода. Памяти он занимает столько же независимо от объёма данных — но лишь пока их обрабатывают по одному: list() сводит эту экономию на нет. Это делает генераторы естественным выбором для больших, потоковых и даже бесконечных последовательностей и неудачным — там, где нужны индексы, длина или повторные проходы.