Итераторы и генераторы в Python за 7 минут

for в Python работает с чем угодно, что умеет отдавать элементы по одному. Это протокол итерации, и на нём построены генераторы, zip, map, чтение файлов и половина стандартной библиотеки.

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

• Итерируемый объект (iterable) умеет выдать итератор: у него есть __iter__. Это list, str, dict, range, файл.
• Итератор (iterator) помнит позицию и выдаёт следующий элемент через __next__, а в конце бросает StopIteration. У итератора __iter__ возвращает его самого.

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

it = iter(items)          # items.__iter__()
while True:
    try:
        x = next(it)      # it.__next__()
    except StopIteration:
        break
    ...

Главное отличие: список можно обойти сколько угодно раз, каждый iter() даёт новый итератор. Итератор одноразовый: дошёл до конца, и он пуст навсегда.

nums = map(int, ['1', '2'])
list(nums)   # [1, 2]
list(nums)   # []  уже исчерпан

Генераторы

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

def countdown(n):
    print('старт')
    while n > 0:
        yield n
        n -= 1

g = countdown(3)    # ничего не напечатано
next(g)             # 'старт', затем 3
list(g)             # [2, 1]

return value в генераторе завершает итерацию, значение попадает в StopIteration.value (его забирает yield from).

Генераторные выражения

(x * x for x in data) то же, что list comprehension, но лениво: элементы вычисляются по требованию, память O(1). Хорошо для sum, any, max и больших файлов:

total = sum(len(line) for line in open('big.log'))

Ленивость и конвейеры

Генераторы складываются в конвейер без промежуточных списков:

lines = (l.strip() for l in f)
errors = (l for l in lines if 'ERROR' in l)
first10 = itertools.islice(errors, 10)

Обрабатывается ровно столько строк, сколько нужно для первых десяти ошибок.

yield from и send

yield from sub отдаёт наружу все элементы sub и возвращает его return-значение. Удобно для рекурсивных обходов:

def walk(tree):
    yield tree.value
    for child in tree.children:
        yield from walk(child)

У генератора есть send(value) (значение станет результатом yield), throw() и close(). Исторически на этом строились корутины, сейчас для асинхронности есть async def.

itertools в двух строках

count, cycle, repeat бесконечные; chain, islice, zip_longest, groupby (группирует соседние равные элементы, данные сначала сортируют), product, permutations, combinations, accumulate, batched (3.12+, пачки по n).

Частые ошибки

• Повторный обход исчерпанного итератора (map, zip, генератора, файла).
• groupby по несортированным данным.
• Ожидание, что тело генератора выполнится при вызове функции.
• Генератор, внутри которого try/finally держит ресурс, а генератор бросили недочитанным: finally выполнится только при close() или сборке мусора.

Что спросят на собесе

• Чем итерируемый объект отличается от итератора, как работает for.
• Что такое генератор, чем yield отличается от return.
• Чем генераторное выражение отличается от list comprehension по памяти.
• Почему map или zip можно обойти только один раз.
• Что делает yield from.