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

Как работает for

Итерируемый объект умеет выдать итератор через __iter__. Итератор выдаёт значения по одному через __next__ и бросает StopIteration, когда они кончились. for делает ровно это:

it = iter([1, 2, 3])
next(it)  # 1
next(it)  # 2

Список итерируемый, но не итератор: его можно обходить сколько угодно раз. Итератор одноразовый, его __iter__ возвращает сам себя — второй проход ничего не даст.

Генераторы

Функция с yield возвращает генератор — итератор, который отдаёт значения по одному и замирает между ними, сохраняя состояние:

def read_lines(path):
    with open(path) as f:
        for line in f:
            if line.strip():
                yield line.rstrip("\n")

Следующее значение вычисляется только по запросу, поэтому память — O(1) вместо целого списка: так обрабатывают файлы на гигабайты и бесконечные потоки.

Списковое включение и генераторное выражение

squares = [x * x for x in data]      # весь список в памяти сразу
total = sum(x * x for x in data)     # лениво, по одному

Список нужен, если результат обходят несколько раз, берут len или индекс. Для однократного прохода в sum, any, max и подобных — генератор.

Срезы

seq[start:stop:step] — новая последовательность с start включительно до stop не включая. Отрицательные индексы считаются с конца, a[::-1] — развёрнутая копия, a[:] — поверхностная копия. Срез не бросает IndexError при выходе за границы — просто вернёт меньше. Для итераторов срезов нет — для них itertools.islice.

send, throw и yield from

Генератор умеет и принимать значения: send(x) передаёт x внутрь, и оно становится результатом выражения yield. throw(exc) бросает исключение в точке yield, close() завершает генератор. Первый вызов — только next() или send(None): генератор ещё не дошёл до yield.

yield from sub() делегирует подгенератору: пробрасывает его значения, send и throw, а возвращает его return. На этих механизмах исторически был построен asyncio до появления async и await. Сейчас send нужен редко, а yield from удобен для композиции генераторов.

itertools

Ленивые строительные блоки:

groupby группирует только подряд идущие одинаковые ключи — без предварительной сортировки один ключ окажется в нескольких группах.

from itertools import batched
for chunk in batched(read_lines("ids.txt"), 500):
    save(chunk)

Что запомнить

• Итерируемый выдаёт итератор, итератор — значения, и он одноразовый.
• Генератор вычисляет по запросу и не держит всё в памяти.
• Однократный проход — генераторное выражение, повторный — список.
• groupby работает по отсортированным данным.