Что такое yield в Python и как его использовать

В Python есть генераторы и итераторы. Это одни из самых часто используемых и основных вещей, когда нам нужно обработать сразу много данных. Но сходу разобраться в итераторах и генераторах может быть сложно, потому что эта тема очень абстрактная. Чтобы было проще вникнуть, мы подробно объясним и расскажем, при чём здесь ключевое слово yield и как оно работает.

Итераторы и их связь с генераторами

В статье мы разберём 3 важных понятия в Python:

Начнём с самого простого.

Итерируемые объекты — всё, что состоит из элементов, которые можно перебирать циклами. Это могут быть строки, списки, словари, кортежи.

На практике нам часто нужно достать какой-то элемент из набора, но итерируемый объект не умеет выдавать свои объекты. Это как книга: если нам нужна какая-то определённая страница, сама по себе книга на ней не откроется.

Итератор — конструкция, которая умеет перебирать элементы итерируемых объектов. Если итерируемый объект можно сравнить с книгой, то итератор — это читатель, который перелистывает её страницы.

Каждый раз, когда мы в программе создаём циклы for для перебора элементов, в этом процессе участвуют итерируемые объекты и итераторы. Мы не видим этого явно, поэтому для нас работа с циклами выглядит так:

Если углубиться в детали, то это будет выглядеть так:

Итераторы — одноразовые. Поэтому при каждом запуске цикла создаётся новый.

Генераторы — это тоже итераторы, и в их исходном коде есть такие же методы, что и в итераторах. Но работают они немного иначе.

Что такое генераторы и как они работают

Генератор — это объект, который позволяет постепенно вычислять и возвращать значения без итерируемого объекта. Он создаётся двумя способами:

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

# создание функции-генератора
def my_generator():
   yield 1
   yield 2
   yield 3

# запускаем генератор
for value in my_generator():
   print(value)

Этот код выведет три числа:

1
2
3

Что такое yield в Python: когда код доходит до yield, он возвращает то, что указано после этого слова, и замораживает выполнение функции. При следующем вызове генератор продолжает работу с этого места. В нашем примере мы последовательно получаем сначала 1, потом 2, потом 3 — по одному результату на каждый вызов функции.

Когда инструкции заканчиваются, генератор выбрасывает исключение StopIteration. Программа понимает, что элементы закончились, и останавливает работу.

Генераторное выражение без явного указания каждого значения занимает меньше места и не использует слово yield:

# генераторное выражение, которое создаёт 
# квадраты натуральных чисел от 0 до 4:
generator = (x ** 2 for x in range(5))

for value in generator:
    print(value)

Код выводит 5 отдельных чисел:

0
1
4
9
16

Этот генератор не использует у себя yield в явном виде, но работает по тому же принципу.

Принцип работы генератора называется ленивым вычислением: он выдаёт по одному значению за раз и помнит, где остановился.

Преимущества использования yield

Генераторы сильно экономят место в памяти. Дело в том, что итератору нужен итерируемый объект, с которым он будет работать. А итерируемому объекту нужно место в памяти для сохранения значений.

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

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

# импортируем модуль, чтобы посчитать объём элементов
import sys

# создаём большой список, который будет храниться в памяти
big_list = [x for x in range(1000000)]

# создаём генератор, который генерирует такую же последовательность
def function():
   for x in range(1000000):
       yield x

# сравниваем размеры в памяти
print(f"Размер списка в памяти: {sys.getsizeof(big_list)} байт")
print(f"Размер генератора в памяти: {sys.getsizeof(function)} байт")

Смотрим, какой результат выходит в консоли запуска:

Размер списка в памяти: 8448728 байт
Размер генератора в памяти: 152 байт

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

Тот же результат с небольшим использованием памяти можно получить, если вместо функции с yield написать генераторное выражение:

big_generator = (x for x in range(1000000))

Недостатки использования yield

Использовать генератор вместо итерируемых объектов не всегда целесообразно:

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

Сравнение yield и return

Основное отличие в том, что return завершает функцию и возвращает одно значение. А yield только приостанавливает выполнение функции и запоминает состояние. При следующем вызове он возвращает следующее значение и снова останавливается, потому что ленивый.

Простой пример — функция всегда останавливается на первом return и никогда не перейдёт дальше:

# функция с return
def value():
   return 1
   # этот код никогда не выполнится:
   return 2
   return 3

print("вызываем функцию первый раз: ", value())
print("вызываем функцию второй раз: ", value())
print("вызываем функцию третий раз: ", value())

Сколько бы раз мы ни вызывали эту функцию, в консоли всегда будет один результат:

вызываем функцию первый раз: 1
вызываем функцию второй раз: 1
вызываем функцию третий раз: 1


А вот аналогичная функция, где return заменён на yield:

# функция с yield
def simple_yield():
   yield 1
   # эти строки выполнятся при следующих вызовах `next()`
   yield 2
   yield 3

# создаём генератор
gen = simple_yield()
print("вызываем генератор первый раз: ", next(gen))
print("вызываем генератор второй раз: ", next(gen))
print("вызываем генератор третий раз: ", next(gen))

Каждый раз генератор переходит на следующую строку с yield:

вызываем генератор первый раз: 1
вызываем генератор второй раз: 2
вызываем генератор третий раз: 3


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

Примеры использования yield в Python

Вот несколько примеров генераторов.

Генератор чисел Фибоначчи. Это бесконечная последовательность чисел, каждое из которых равно сумме двух предыдущих.

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

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

# функция-генератор чисел Фибоначчи
def fibonacci():
   # задаём начальные числа последовательности:
   a, b = 0, 1
   # создаём бесконечный цикл
   while True:
       # возвращаем текущее значение `a`
       yield a
       # обновляем `a` и `b`
       a, b = b, a + b

# сохраняем функцию в отдельную переменную
keyword = fibonacci()
# вызываем генератор несколько раз
print(next(keyword))
print(next(keyword))
print(next(keyword))
print(next(keyword))

Проверяем работу:

0
1
1
2


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

В этом примере мы работаем с файлом, в котором 100 000 строк:

# создаём функцию для чтения файлов
def read_file(filename):
   # создаём контекст-менеджер для открытия файла
   with open(filename) as f:
       # для каждой строки в файле...
       for line in f:
           # возвращаем строку без лишних пробелов
           yield line.strip()

# создаём цикл для генератора
for line in read_file("big.txt"):
   # на один yield генератор выводит одну строку
   # без загрузки всего файла в память
   print(line)

При запуске в консоли строки файла выводятся по одной:

Файл для такого же теста можно создать через такую программу. Если что, в Python две звёздочки подряд — это возведение в степень:

with open("big.txt", "w") as f:
   for i in range(10**6):
       f.write(f"Строка {i}\n")

Генератор простых чисел. Простые числа имеют всего два делителя: единицу и само это число. Код создаёт бесконечную последовательность таких чисел. Основная идея — проверять каждое число n, чтобы оно делилось только на 1 и само себя. Если да, оно возвращается через yield.

# создаём генератор простых чисел
def primes():
   # создаём начальную точку генерации
   # первое натуральное простое число — 2
   n = 2
   # запускаем бесконечный цикл
   while True:
       # проверяем, что число не делится без остатка
       # на какие-то другие числа в диапазоне от 2 до квадратного корня из числа
       if all(n % d != 0 for d in range(2, int(n**0.5) + 1)):
           # если `n` простое, возвращаем его
           yield n
       # после проверки переходим к следующему числу
       n += 1

# сохраняем генератор в переменную
g = primes()
# запускаем генератор 5 раз
for _ in range(5):
   print(next(g))

Проверяем, что получается при запуске:

2
3
5
7
11


Основные принципы работы с yield

Вот что нужно запомнить при использовании генераторов и yield:

Контроль за исчерпанием генератора

Если генератор закончил считать значения, при следующем вызове он выдаёт исключение — ошибку StopIteration.

Самый простой и удобный способ работы с генераторами — через цикл for. Он автоматически останавливается при исчерпании оператора.

Другой вариант — обрабатывать исключение StopIteration через конструкцию try-except. Идея в том, что если в генераторе больше не осталось значений, то он выдаст ошибку, в ответ на которую можно запрограммировать какие-то действия. Мы в примере просто выведем на экран сообщение, но можно придумать что-то более сложное:

# создаём функцию-генератор
def gen():
   # возвращаем первое значение
   yield 1
   # возвращаем второе значение
   yield 2

# сохраняем генератор в переменную
values = gen()

# пытаемся вызвать генератор
try:
   # вызываем первый раз
   print(next(values))
   # вызываем второй раз
   print(next(values))
   # вызываем третий раз
   print(next(values))
# обрабатываем исключение StopIteration
except StopIteration:
   print("Генератор закончился!")

Теперь при вызове программа будет работать так:

1
2
Генератор закончился!


Также можно просто создать ещё один генератор:

# создаём функцию-генератор
def gen():
   yield from range(3)

# сохраняем генератор в переменную
used = gen()
# делаем ещё одну копию генератора
generators = gen()

print("Выводим значения генератора statement:", list(used))
print("Выводим значения генератора generators:", list(generators))

Второй генератор выводит такие же значения, как первый:

Выводим значения генератора statement: [0, 1, 2]
Выводим значения генератора generators: [0, 1, 2]


Практические советы и лучшие практики

При работе с yield есть несколько вещей, которые стоит помнить и которыми нужно пользоваться: