Итераторы и генераторы в Python: как экономить память при обработке данных

Когда программа обрабатывает большой файл, ей не всегда нужен список всех строк. Часто достаточно взять очередную запись, проверить её и сразу включить в расчёт. Итераторы и генераторы позволяют организовать такую обработку по требованию. Экономия появляется за счёт отказа от лишних накопленных коллекций, а не из-за самого слова yield в коде.

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

Iterable и iterator: контейнер и текущая позиция

Итерируемый объект умеет предоставить итератор. Например, список можно передать iter(), получить объект с текущей позицией и запрашивать у него значения через next(). Сам список и этот итератор — разные объекты:

numbers = [3, -2, 7]
first = iter(numbers)
print(next(first))  # 3
print(next(first))  # -2
second = iter(numbers)
print(next(second))  # снова 3

Здесь два независимых прохода по одному списку. Итератор, напротив, обычно возвращает самого себя из своего __iter__(), а его __next__() выдаёт следующий элемент. Когда элементов больше нет, он сообщает об этом исключением StopIteration. Цикл for использует этот протокол и завершает перебор без необходимости писать обработчик вручную.

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

Когда нужен явный next()

Для обычной обработки всех значений используйте for. Явный next() нужен, когда вы намеренно забираете отдельный элемент: например, смотрите первую запись, а оставшиеся обрабатываете циклом. Такой просмотр сдвигает позицию, поэтому просмотренная запись уже не попадёт в последующий перебор того же итератора.

items = iter(["a", "b"])
print(next(items, None))  # a
print(list(items))        # ['b']
print(next(items, None))  # None

Второй аргумент next() задаёт результат при исчерпании вместо StopIteration. Но если None может быть настоящим элементом, он не отличает конец потока от такого значения. Для различения можно использовать отдельный объект-маркер и сравнивать по идентичности. Не выбирайте специальное значение вроде нуля, если ноль допустим в данных.

Свой итератор полезно один раз написать ради понимания состояния:

class Countdown:
    def __init__(self, start):
        self.remaining = start

    def __iter__(self):
        return self

    def __next__(self):
        if self.remaining <= 0:
            raise StopIteration
        value = self.remaining
        self.remaining -= 1
        return value

print(list(Countdown(3)))  # [3, 2, 1]

Это учебная реализация для неотрицательного целого начала. Проверка входного типа здесь опущена, чтобы было видно главное: объект хранит позицию между вызовами. Каждый новый Countdown(3) создаёт новый проход, а повторный list() на уже исчерпанном экземпляре не восстановит значения.

yield сохраняет место продолжения

Генераторная функция описывает похожее поведение короче:

def countdown(start):
    while start > 0:
        yield start
        start -= 1

steps = countdown(3)
print(next(steps))  # 3
print(next(steps))  # 2
print(list(steps))  # [1]

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

Обычный return завершает функцию; в генераторе он также завершает выдачу. Не используйте raise StopIteration внутри генераторной функции как способ закончить поток: для этого достаточно return или естественного конца тела. Протокол завершения при использовании генератора обеспечивает Python.

Выражение-генератор или список

Списковое включение собирает результат сразу. Генераторное выражение выдаёт элементы по запросу:

numbers = [3, -2, 7]
ready = [x * x for x in numbers if x > 0]
lazy = (x * x for x in numbers if x > 0)
print(ready)       # [9, 49]
print(next(lazy))  # 9
print(sum(lazy))   # 49: первое значение уже забрали

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

В примере исходный numbers всё равно остаётся в памяти. Мы избавились только от дополнительного списка квадратов. Если затем вызвать list(lazy), оставшиеся результаты вновь будут собраны в коллекцию. Так же отдельного анализа требуют промежуточные словари, множество уникальных значений и накопление строк: ленивый вход не делает их бесплатными по памяти.

Конвейер чтения файла

Создайте учебный amounts.txt с тремя строками: 3, -2 и 7. Следующая функция принимает любой источник текстовых строк, пропускает пустые и выдаёт положительные целые:

def positive_amounts(lines):
    for line in lines:
        text = line.strip()
        if not text:
            continue
        value = int(text)
        if value > 0:
            yield value

with open('amounts.txt', encoding='utf-8') as stream:
    amounts = positive_amounts(stream)
    total = sum(amounts)
print(total)  # 10

Файловый объект перебирается по строкам, поэтому мы не вызываем readlines() и не создаём список всего содержимого. Генератор хранит состояние обработки, а sum потребляет значения внутри with, пока файл открыт. Если в строке будет «ошибка», int() вызовет ValueError при потреблении потока. Это ожидаемое ограничение нашего формата, а не ситуация, которую нужно молча превращать в ноль.

Не выносите потребление наружу так: создали positive_amounts(stream) внутри with, вышли из блока, а sum вызвали потом. Генератор ещё не прочитал строки, но файл уже закрыт. Ленивая обработка переносит во времени не только вычисления, но и возможные ошибки и требования к доступности ресурса.

Для следующих преобразований можно добавить генераторное выражение:

with open('amounts.txt', encoding='utf-8') as stream:
    squares = (x * x for x in positive_amounts(stream))
    total_squares = sum(squares)
print(total_squares)  # 58

Это авторская цепочка из чтения, преобразования текста, фильтра и расчёта. Она не хранит все строки или все квадраты одновременно. Однако одна огромная строка, большое полученное значение и сам исходный ресурс всё равно потребляют память. Поэтому корректнее говорить об отсутствии накопления всей последовательности, а не о нулевом или неизменно одинаковом расходе.

Как не потерять данные при повторном проходе

Частая ошибка — сначала вывести list(amounts) ради отладки, а потом ожидать прежний результат от sum(amounts). Первый вызов уже исчерпал поток. Для повторения заново откройте файл и создайте генератор либо осознанно сохраните небольшой результат в список. Выбор зависит от объёма и дальнейших действий.

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