2.31. Итераторы и генераторы

Цикл for выполняет больше работы, чем кажется. Эта страница описывает протокол итератора, на котором он работает, и ключевое слово yield, позволяющее создавать собственные итераторы.

2.31.1. Протокол итератора

Каждый объект, который можно перебрать в цикле, реализует два метода:

• __iter__() – возвращает итератор по элементам объекта.

• __next__() – у итератора возвращает следующий элемент или возбуждает StopIteration, когда элементов больше нет.

Встроенная функция iter() вызывает __iter__; next() вызывает __next__. Пройдём по списку вручную:

it = iter([10, 20, 30])
print(next(it))    # 10
print(next(it))    # 20
print(next(it))    # 30
print(next(it))    # raises StopIteration

for – это синтаксический сахар для «вызвать __iter__ один раз, затем выполнять __next__ в цикле, пока не возникнет StopIteration».

Что на самом деле делает for x in items::

_it = iter(items)
while True:
    try:
        x = next(_it)
    except StopIteration:
        break
    # ... loop body ...

Каждый список, кортеж, строка, словарь, множество, файловый объект и генератор уже реализуют __iter__ и __next__ – именно поэтому все они работают с for.

2.31.2. yield и функции-генераторы

Функция, содержащая инструкцию yield, является функцией-генератором. Её вызов не запускает тело; он возвращает объект-генератор (итератор), который выполняет тело по одному yield за раз:

def count_up_to(n):
    i = 0
    while i < n:
        yield i
        i += 1

for value in count_up_to(3):
    print(value)

Вывод:

0
1
2

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

next() выполняет тело до следующего yield, передаёт значение обратно и приостанавливается. Локальное состояние переживает приостановку.

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

2.31.3. Ленивые конвейеры

Генераторы хорошо компонуются. Вывод одного генератора может подаваться на вход другому:

def numbers():
    i = 0
    while True:
        yield i
        i += 1

def squares(source):
    for x in source:
        yield x * x

pipeline = squares(numbers())

for v in pipeline:
    if v > 100:
        break
    print(v)

Значения проходят через конвейер по одному за раз – без промежуточного списка, без встроенной верхней границы в numbers, и потребитель (for v in pipeline) решает, когда остановиться.

Каждый next() у потребителя запускает одно извлечение по всей цепочке; значения существуют только тогда, когда что-то их запрашивает.

2.31.3.1. yield from

Цикл, который извлекает элементы из другого итерируемого объекта и выдаёт каждый из них, встречается достаточно часто, чтобы Python предоставил для него сокращение. Выражение yield from iter выдаёт каждое значение, которое производит итерируемый объект, по порядку – как если бы в генераторе был встроенный цикл for x in iter: yield x:

def chain(*sources):
    for source in sources:
        yield from source

for v in chain([1, 2, 3], (4, 5), "abc"):
    print(v)

Вывод:

1
2
3
4
5
a
b
c

yield from в точности эквивалентен явному циклу for, только короче, и он корректно распространяет StopIteration из внутреннего итерируемого объекта во внешний генератор – что полезно при последовательном связывании нескольких генераторов.

2.31.3.2. Когда yield исчерпан

Достижение конца функции-генератора (или явный return) автоматически возбуждает StopIteration. Нет необходимости возбуждать его вручную; окружающий цикл for видит его и завершается.

Используйте генераторы, когда производящий код естественным образом записывается как цикл с несколькими точками yield; используйте обычное списковое включение, когда вам действительно нужна вся последовательность в памяти.