Генераторы Python: ленивые вычисления и экономия памяти

Проблема: огромные списки

Допустим, нужно прочитать файл на 10 ГБ построчно:

# Плохо — загружает весь файл в RAM
lines = open("huge.log").readlines()
for line in lines:
    process(line)

# Хорошо — генератор
def read_lines(path):
    with open(path) as f:
        for line in f:
            yield line.strip()

for line in read_lines("huge.log"):
    process(line)

Разница: первый вариант занимает гигабайты RAM, второй — несколько килобайт.

Как работает yield?

Ключевое слово yield превращает функцию в генераторную функцию. При вызове такой функции возвращается объект-генератор — он не выполняет тело функции немедленно.

def countdown(n: int):
    print("Старт!")
    while n > 0:
        yield n
        n -= 1
    print("Конец!")

gen = countdown(3)  # ничего не выполнено
next(gen)           # "Старт!", возвращает 3
next(gen)           # возвращает 2
next(gen)           # возвращает 1
next(gen)           # "Конец!", бросает StopIteration

Генератор «замораживается» на каждом yield и продолжает выполнение при следующем next().

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

Аналог list comprehension, но ленивый:

squares_list = [x**2 for x in range(1_000_000)]   # список в памяти
squares_gen  = (x**2 for x in range(1_000_000))    # генератор

import sys
print(sys.getsizeof(squares_list))  # ~8 MB
print(sys.getsizeof(squares_gen))   # ~200 байт

Цепочка генераторов

Генераторы можно соединять в пайплайны:

def integers():
    n = 1
    while True:
        yield n
        n += 1

def take(n, it):
    for _ in range(n):
        yield next(it)

def even(it):
    for x in it:
        if x % 2 == 0:
            yield x

# Первые 5 чётных чисел
result = list(take(5, even(integers())))
# [2, 4, 6, 8, 10]

yield from

yield from делегирует итерацию другому итерируемому объекту:

def flatten(nested):
    for item in nested:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

list(flatten([1, [2, [3, 4]], 5]))
# [1, 2, 3, 4, 5]

Двусторонняя связь: send()

Генераторы могут принимать значения через send():

def accumulator():
    total = 0
    while True:
        value = yield total
        if value is None:
            break
        total += value

gen = accumulator()
next(gen)       # запускаем генератор, получаем 0
gen.send(10)    # 10
gen.send(5)     # 15
gen.send(3)     # 18

Генераторы vs list comprehensions

itertools — батарейки в комплекте

from itertools import islice, chain, groupby

# Первые 10 элементов бесконечного генератора
list(islice(integers(), 10))

# Объединение нескольких итерируемых
list(chain([1, 2], [3, 4], [5]))  # [1, 2, 3, 4, 5]

# Группировка (требует предварительной сортировки)
data = sorted([("a", 1), ("b", 2), ("a", 3)], key=lambda x: x[0])
for key, group in groupby(data, key=lambda x: x[0]):
    print(key, list(group))

Итог

Используйте генераторы когда:

• Данных слишком много для хранения в памяти
• Нужна ленивая оценка (вычисляем только то, что нужно)
• Строите пайплайны обработки данных
• Работаете с бесконечными последовательностями