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

«Работали ли с генераторами?» первый вопрос на большинстве собеседований инженеров данных, и ответ «да, там yield» засчитывают редко. Интервьюер хочет услышать, зачем они нужны: чтобы обрабатывать данные, не загружая их целиком в память, и собирать обработку в конвейер. В уроке про память было видно, как это работает на уровне ссылок. Здесь разбираем сам механизм.

Итерируемое и итератор

Два разных понятия, которые часто смешивают.

Итерируемый объект (iterable) это то, по чему можно пройти циклом for: список, строка, словарь, файл, генератор. Технически у него есть метод __iter__, который возвращает итератор.

Итератор (iterator) это объект с методом __next__, который отдаёт следующий элемент, а когда элементы кончились, возбуждает StopIteration. Итератор помнит, где остановился.

values = [10, 20, 30]
it = iter(values)        # values.__iter__()
print(next(it))          # 10
print(next(it))          # 20
print(list(it))          # [30]: забрали остаток
print(list(it))          # []: итератор исчерпан

Цикл for делает ровно это: вызывает iter() один раз, потом next() до StopIteration. Список можно обойти сколько угодно раз, потому что каждый iter(values) создаёт новый итератор. Итератор обходится один раз: второй list(it) пустой. Это главный источник ошибок с генераторами: передали генератор в функцию, которая его прошла, а потом удивляемся, что в следующем цикле пусто.

Сам итератор тоже итерируем (__iter__ возвращает его самого), поэтому его можно поставить в for.

Файловый объект это итератор по строкам, zip, map, filter, enumerate, reversed, csv.reader возвращают итераторы. Все они ленивые: ничего не вычисляют, пока не попросят следующий элемент.

yield: функция, которая ставится на паузу

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

def numbers():
    print("старт")
    yield 1
    print("между")
    yield 2
    print("конец")
 
gen = numbers()          # ничего не напечатано
print(next(gen))         # старт, 1
print(next(gen))         # между, 2
print(next(gen))         # конец, затем StopIteration

return в генераторе завершает его: дальше StopIteration. Значение после return обычный for не видит.

Разница между return и yield в одном: return отдаёт результат и уничтожает кадр функции, yield отдаёт значение и сохраняет кадр до следующего запроса. Поэтому функция с return должна собрать весь результат до того, как отдать его, а функция с yield отдаёт по одному элементу и ничего не копит.

def read_rows_list(path):
    rows = []
    for line in open(path):
        rows.append(parse(line))
    return rows               # весь файл в памяти
 
def read_rows(path):
    with open(path) as f:
        for line in f:
            yield parse(line) # одна строка в памяти

Второй вариант ещё и закрывает файл правильно: with внутри генератора отработает, когда генератор дочитает файл или будет уничтожен.

Три способа создать генератор

Функция с yield, как выше. Генераторное выражение, как списковое включение, но в круглых скобках:

amounts = (row["amount"] for row in rows if row["status"] == "paid")
total = sum(amounts)

Квадратные скобки создали бы список всех сумм в памяти, круглые отдают по одной. Если результат сразу уходит в sum, max, any, ",".join или в запись в файл, список не нужен.

И yield from, который делегирует другому итерируемому: отдаёт все его элементы, как будто они написаны через yield по одному.

def read_many(paths):
    for path in paths:
        yield from read_rows(path)

Конвейер генераторов

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

def lines(path):
    with open(path, encoding="utf-8") as f:
        yield from f
 
def parsed(lines):
    for line in lines:
        yield line.rstrip("\n").split(",")
 
def only_paid(rows):
    for row in rows:
        if row[1] == "paid":
            yield row
 
def write(rows, out_path):
    with open(out_path, "w", encoding="utf-8") as out:
        for row in rows:
            out.write(",".join(row) + "\n")
 
write(only_paid(parsed(lines("orders.csv"))), "paid.csv")

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

Для готовых кирпичиков есть itertools: islice берёт первые n элементов без загрузки остального (посмотреть начало большого файла), chain склеивает несколько итераторов, groupby группирует соседние элементы с одинаковым ключом (данные должны быть отсортированы), batched в Python 3.12 режет на чанки.

Ленивость переносит ошибки

Генератор не выполняется при создании, поэтому ошибка внутри него возникает не там, где его создали, а там, где его начали читать:

def parse_all(lines):
    for line in lines:
        yield int(line)
 
gen = parse_all(["1", "x", "3"])   # ошибки нет
first = next(gen)                  # 1
second = next(gen)                 # ValueError здесь

Если обернуть в try только создание генератора, исключение пройдёт мимо. Обработчик нужен вокруг потребления: вокруг цикла или внутри самого генератора, если политика «пропустить плохую строку» принадлежит ему.

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

Генератор или список

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

Если нужен и генератор, и повторный проход, материализуйте явно: rows = list(read_rows(path)). Тогда в коде видно, где данные попали в память.

Как отвечать на собеседовании

Что такое итерация и итератор? Итерация это последовательный обход элементов. Итератор это объект с __next__, который отдаёт элементы по одному и возбуждает StopIteration в конце; итерируемый объект это тот, у которого __iter__ возвращает итератор. for вызывает iter, а потом next до конца.

Что такое генератор и как его определить? Итератор, который создаётся функцией с yield или генераторным выражением в круглых скобках; третий способ yield from для делегирования. Тело выполняется лениво, по одному элементу на запрос, и сохраняет состояние между запросами.

В чём разница между return и yield? return завершает функцию и отдаёт один результат, кадр уничтожается. yield отдаёт значение и приостанавливает функцию, сохраняя локальные переменные; при следующем next выполнение продолжается. Функция с yield возвращает генератор, а не значение.

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

Практика