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
A for loop calls __iter__ once to get an iterator, then calls __next__ repeatedly until StopIteration ends the loop.

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 у цьому випадку) зберігається між відновленнями.

Sequence diagram with two lifelines (caller and generator body). The caller calls count_up_to(3), which creates a generator without running the body. Each subsequent next() runs the body until the next yield, returns the yielded value, and pauses. The fourth next() falls off the end and raises StopIteration. The variable i is preserved across pauses.

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) вирішує, коли зупинитися.

Three boxes left-to-right: numbers(), squares(source), and the for-v-in-pipeline consumer. Three cycles are drawn beneath. In each cycle, the consumer sends a pull request leftward to squares, which sends a pull leftward to numbers; numbers yields a value rightward to squares, which yields its squared value rightward to the consumer.

Кожен 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; використовуйте звичайне розуміння списку, коли вам справді потрібна вся послідовність у пам’яті.