Генераторы

Вернёмся к книге с закладкой. Представьте чтеца, который читает вам вслух по одной главе. Дочитав главу, он кладёт закладку и замолкает. Попросите ещё — он откроет книгу на закладке и прочтёт следующую главу, а его пометки на полях останутся на месте.

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

def evens(limit):
    for n in range(0, limit, 2):
        yield n


g = evens(7)
print(next(g), next(g))  # 0 2
print(list(g))  # [4, 6]

Вызов evens(7) не выполняет тело, а возвращает генератор — итератор, который вычисляет значения по запросу. Давайте разберём пример по строкам:

1. g = evens(7) создаёт генератор, но ни одна строка тела ещё не выполнилась.
2. Первый next(g) выполняет тело до yield n и получает 0.
3. Второй next(g) продолжает с места остановки и получает 2.
4. list(g) забирает оставшиеся значения: [4, 6].

Пауза на yield

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

Нажимайте «Следующий шаг» и следите за рамкой counter. Она появляется на каждом next и уходит на yield, а n при этом не сбрасывается в 1.

Когда тело заканчивается, генератор вызывает StopIteration, и for, list или sum завершаются без ошибки. return в теле генератора тоже завершает его досрочно.

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

Один проход

Генератор — итератор, поэтому значения он отдаёт один раз. Здесь аналогия с чтецом заканчивается: чтец может начать книгу заново, а генератор назад не возвращается. Возьмём g = evens(5):

• первый list(g) даёт [0, 2, 4];
• второй list(g) даёт [];
• list(evens(5)) снова даёт [0, 2, 4], потому что это новый генератор.

Второй цикл for по тому же g не выполнит тело ни разу, а sum(g) вернёт 0, как от пустого списка. Чтобы пройти значения ещё раз, вызовите функцию заново или сохраните результат в список. Генераторное выражение (n * n for n in nums) из темы о включениях создаёт такой же одноразовый генератор.

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

Генератор считает значение только тогда, когда его просят. Это то же ленивое вычисление, что и у генераторного выражения:

def countdown(n):
    print("start", end=" ")
    while n > 0:
        yield n
        n -= 1


g = countdown(3)
print("created")  # created
print(next(g))  # start 3

Слово start появилось только при первом next, а не при вызове countdown(3). Поэтому генератор может быть и бесконечным, например с while True: в теле: значения будут вычисляться, пока их забирают. Цикл for по такому генератору останавливают через break, и оставшиеся значения так и не вычисляются.

Частая ошибка

У генератора спрашивают длину:

g = evens(7)
print(len(g))  # ✗

TypeError: object of type 'generator' has no len()

Генератор не знает, сколько значений выдаст, пока не дойдёт до конца тела. Если нужна длина или индексы, соберите значения в список: len(list(evens(7))) даёт 4.

Главное

• Функция с yield при вызове возвращает генератор и не выполняет тело.
• Каждый next выполняет тело до следующего yield и ставит генератор на паузу, сохраняя локальные переменные.
• Генератор проходится один раз, а значения считает только по запросу.
• У генератора нет len: сначала соберите значения в список.