Генераторы в Python

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

PythonКопировать

def sum(val, m):
    result = val
    while True:
        print(result)
        result += m

print(sum(0, 1))
print("Конец программы")

Если запустить программу, то функция будет работать бесконечно, так как в примере не предусмотрено завершение цикла. Весь код, следующий за вызовом sum, никогда не будет выполнен.

Изменим код:

PythonКопировать

def sum(val, m):
    result = val
    while True:
        # заменили print на yield
        yield result
        result += m

print(sum(0, 1))
print("Конец программы")

# output
# >>> <generator object sum at 0x7f1032f95a50>
# >>> Конец программы

Вызов функции sum вычислился в <generator object>, поэтому она не зациклилась.

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

PythonКопировать

for number in sum(0, 3):
    print(number)
    if number > 11:
        break

# output
# >>> 0
# >>> 3
# >>> 6
# >>> 9
# >>> 12

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

Посмотрим, как будет выглядеть тот же самый код, если написать его с использованием генератора.

PythonКопировать

def gen_range(stop_value):
    stop_value = stop_value - 1
    current = -1
    while current < stop_value:
        current += 1
        yield current


for x in gen_range(3):
    print(x)

# output
# >>> 0
# >>> 1
# >>> 2

Работа генераторов построена на принципе запоминания контекста выполнения функции. Функция-генератор запоминает, на каком месте она остановилась, и может продолжить выполнение после ключевого слова yield.

Рассмотрим, как поведёт себя генератор с несколькими yield:

PythonКопировать

def simple_generator():
    yield 1
    yield 2
    return 3

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

PythonКопировать

gen = simple_generator()
print(next(gen))
print(next(gen))
print(next(gen))

В результате вы увидите следующее:

PythonКопировать

1
2
Traceback (most recent call last):
...
StopIteration: 3

То есть функция действительно запоминает, где она остановилась после каждого вызова функции next.

Вопрос о разнице между итераторами и генераторами в Python довольно частый и актуальный. Обе сущности сильно связаны, и их часто путают, что приводит к недопониманиям.

Итератор — более общая концепция. Это объект, у которого определены два метода: __next__ и __iter__. С другой стороны, генератор — это итератор, но не наоборот. Если посмотреть вывод dir(<generator>), то обнаружим, что в нём тоже определены методы __next__ и __iter__. Также важно помнить, что из генераторов нельзя получить значения второй раз.

Применение генераторов

Из-за своей «ленивой» обработки, когда данные обрабатываются не все сразу, а частями, генераторы нашли широкое применение. Разберём наиболее распространённые кейсы, когда могут пригодиться генераторы.

Бесконечные последовательности

Генератор позволяет создавать последовательности бесконечной длины.

PythonКопировать

def func(x):
    n = x
    while True:
        n = (n + 3) * (n + 2) - 5 * n
        yield n
        if n > 100:
            break

for item in func(2):
    print(item)

# output
# >>> 10
# >>> 106

Чтение больших файлов

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

PythonКопировать

def read_file(file_name):
    with open(file_name) as f:
        return f.read().split("\n")

read_my_file = read_file("my_file.txt")

row_count = 0
for row in read_my_file:
    row_count += 1

print(f"В файле {row_count} строк.")

Если размер файла окажется слишком большим, то данная программа выполнится с ошибкой. Чтобы избежать этого, используем генератор:

PythonКопировать

def read_file(file_name):
    with open(file_name) as f:
        for row in f:
            yield row