Итераторы, генераторы и корутины

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

Эти механизмы позволяют обрабатывать файлы, не помещающиеся в память, и описывать бесконечные последовательности; на генераторах построена вся асинхронность Python — asyncio рассматривается в главе, посвящённой асинхронности.

Итераторы

Любая коллекция, встроенная в язык (списки, словари, множества, строки, файлы), является итерируемой.

Итератор представляет собой объект с методом __next__, который при каждом вызове возвращает следующий элемент, а по исчерпании элементов возбуждает StopIteration. Итерируемый объект обладает методом __iter__, возвращающим итератор. Список является итерируемым, но не итератором, поэтому по нему можно пройти дважды, а по итератору — лишь один раз.

Рассмотрим собственный аналог range, в котором обе роли разделены явно. Контейнер Range способен только выдать итератор, а сам обход осуществляется классом RangeIterator.

class Range:
    def __init__(self, stop_value: int):
        self.stop_value = stop_value - 1

    def __iter__(self):
        return RangeIterator(self)

class RangeIterator:
    def __init__(self, container):
        self.container = container
        self.current = -1          # состояние обхода живёт здесь

    def __iter__(self):
        return self                # итератор обязан быть итерируемым

    def __next__(self):
        if self.current < self.container.stop_value:
            self.current += 1
            return self.current
        raise StopIteration

Поле current хранится в итераторе, а не в контейнере, поэтому каждый новый for начинает счёт заново, а по одному объекту Range(5) можно пройти произвольное число раз. Если поместить счётчик в контейнер, второй проход окажется пустым. Без метода __iter__ у самого итератора объект не является итерируемым, и цикл for по нему невозможен.

Чаще обе роли совмещают в одном классе, где __iter__ возвращает self: кода при этом меньше, однако второй проход начнётся там, где закончился первый.

class Range2:
    def __init__(self, stop_value: int):
        self.current = -1
        self.stop_value = stop_value - 1
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current < self.stop_value:
            self.current += 1
            return self.current
        raise StopIteration

Обычный цикл for получает итератор через iter, вызывает next в цикле и перехватывает StopIteration, чтобы остановиться.

iterable = Range2(5)
iterator = iter(iterable)

while True:
    try:
        value = next(iterator)
        print(value)
    except StopIteration:
        break

Генераторы

Генераторы работают на принципе запоминания контекста, сохраняемого ключевым словом yield.

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

В примере ниже первые два вызова возвращают значения, а третий доходит до return и возбуждает StopIteration, поместив в него возвращённое значение.

def simple_generator():
    yield 1
    yield 2
    return 3

gen = simple_generator()
print(next(gen))  # 1
print(next(gen))  # 2
print(next(gen))  # StopIteration: 3

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

gen_exp = (x for x in range(100000))
print(gen_exp)  # <generator object <genexpr> at 0x...>

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

def read_events(path):
    """Отдаёт события журнала по одному, не читая файл целиком."""
    with open(path) as f:
        for line in f:
            if not line.strip() or line.startswith('#'):
                continue           # пустые строки и комментарии пропускаем
            t, channel, value = line.split()
            yield float(t), channel, float(value)

for t, channel, value in read_events('run042.log'):
    if channel == 'BPM01':
        print(t, value)

Обратный случай — бесконечная последовательность. Генератор вычисляет значения, пока они запрашиваются, а границу устанавливает читающая сторона: itertools.islice берёт первые несколько значений и освобождает генератор.

import itertools

def clock(dt):
    """Бесконечная последовательность моментов времени с шагом dt."""
    t = 0.0
    while True:
        yield t
        t += dt

for t in itertools.islice(clock(0.02), 5):
    print(round(t, 3))  # 0.0, 0.02, 0.04, 0.06, 0.08 — по одному в строке

Когда генератор лишь передаёт элементы другого источника, цикл с yield записывается короче через yield from. Это не только синтаксическое упрощение: конструкция пробрасывает во вложенный генератор всё, что передаётся через send и throw, а наружу отдаёт его return, что необходимо при вложенных генераторах.

numbers = [1, 2, 3]

# Стандартный подход
def func():
    for item in numbers:
        yield item

# Упрощенный подход
def func():
    yield from numbers

Корутины

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

Генератор только отдаёт значения наружу, а корутина способна и принимать их. Выражение deposit = (yield) приостанавливает функцию и ожидает, пока извне будет вызван send. Таким образом, получается диалог с функцией, сохраняющей между запросами своё состояние.

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

import math
from collections.abc import Generator

def cash_return_coro(percent: float,
                     years: int) -> Generator[float | None, float, None]:
    value = math.pow(1 + percent / 100, years)
    while True:
        try:
            deposit = (yield)
            yield round(deposit * value, 2)
        except GeneratorExit:
            print('Выход из корутины')
            raise

# Использование
coro = cash_return_coro(5, 5)
next(coro)
values = [1000, 2000, 5000, 10000, 100000]
for item in values:
    print(coro.send(item))
    next(coro)
coro.close()

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

Дальнейшее изучение

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

На этом построена асинхронность в Python: цикл событий представляет собой планировщик, который переключается между приостановленными генераторами и отдаёт процессорное время той задаче, которая готова работать, пока остальные ожидают ответа от сети или диска. Внутреннее устройство этого механизма и его использование через asyncio рассматриваются в главе про асинхронность в разделе «Ускорение расчётов».