ЛУЧШЕЕ ПРОГРАММИРОВАНИЕ

Как использовать генератор и yield в Python

Работа с большими наборами данных или файлами с помощью генераторов Python

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

Что такое генераторы в Python?

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

Функции генератора, представленные в PEP 255, представляют собой особый вид функций, которые возвращают своего рода ленивый итератор. Есть объекты, которые можно перебирать как список, однако, в отличие от списков, ленивые итераторы не хранят свое содержимое в памяти. Одним из преимуществ функций генератора перед итераторами является объем кода, необходимый для кодирования.

После этого введения давайте посмотрим на несколько примеров генераторов в действии:

Некоторые варианты использования генераторов

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

Обычный вариант использования генераторов - это работа с большими файлами или потоками данных, например, с файлами CSV. Допустим, нам нужно посчитать, сколько строк находится в текстовом файле, наш код может выглядеть примерно так:

csv_gen = csv_reader("some_file.txt")
row_count = 0

for row in csv_gen:
    row_count += 1

print(f"Row count is {row_count}")

с нашей функцией csv_reader, реализованной следующим образом:

def csv_reader(file_name):
    file = open(file_name)
    result = file.read().split("\n")
    return result

Теперь это очень ясно и просто. Наша функция csv_reader просто откроет файл в памяти и прочитает все строки, затем она разделит строки и сформирует массив с данными файла, поэтому наш приведенный выше код будет работать идеально, по крайней мере, мы так думаем.

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

Traceback (most recent call last):
  File "ex1_naive.py", line 22, in <module>
    main()
  File "ex1_naive.py", line 13, in main
    csv_gen = csv_reader("file.txt")
  File "ex1_naive.py", line 6, in csv_reader
    result = file.read().split("\n")
MemoryError

Мы вылетели из программы. Файл был слишком большим и не мог быть загружен в память, из-за чего python генерировал исключение и аварийно завершал работу.

Итак, как мы можем это исправить? ну ... мы знаем, что с помощью генераторов у нас есть способ создавать простые итераторы, так что это должно помочь, давайте теперь взглянем на функцию csv_reader, созданную с помощью генератора.

def csv_reader(file_name):
    for row in open(file_name, "r"):
        yield row

Все еще очень просто, это выглядит еще элегантнее, чем раньше, но что это за ключевое слово yield?

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

Когда функция содержит yield, Python автоматически (и за кулисами) реализует итератор, применяя для нас все необходимые методы, такие как __iter__() и __next__(), поэтому нам не нужно ни о чем беспокоиться.

Возвращаясь к нашему примеру, если мы теперь решим выполнить наш код, мы получим следующее:

Row count is 65123455

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

Но это еще не конец истории, есть еще более простые и интересные способы реализации генераторов путем определения выражения генератора (также называемого понимание генератора), имеющего синтаксис это очень похоже на понимание списка.

Посмотрим, как это будет выглядеть

csv_gen = (row for row in open(file_name))

Красиво, не правда ли? просто запомните эти основные отличия:

  • Использование yield приведет к созданию объекта-генератора
  • Использование return приведет к отображению только первой строки файла.

Создание бесконечной последовательности

Другой распространенный сценарий для генераторов - генерация бесконечной последовательности. В Python, когда вы используете конечную последовательность, вы можете просто вызвать range() и оценить ее в контексте списка, например:

a = range(5)
print(list(a))
[0, 1, 2, 3, 4]

Мы могли бы сделать то же самое, сгенерировав бесконечную последовательность с помощью таких генераторов:

def infinite_sequence():
    num = 0
    while True:
        yield num
        num += 1

И вы можете использовать его, например, для печати значений

for i in infinite_sequence():
    print(i, end=" ")

Хотя это будет очень быстро и будет работать «вечно», поэтому вам придется остановить его вручную, нажав CTRL+C или альтернативу MAC, но вы увидите, что все числа будут напечатаны очень быстро на экране.

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

>> gen = infinite_sequence()
>>> next(gen)
0
>>> next(gen)
1
>>> next(gen)
2
....

Подробнее об уступке

До сих пор мы рассматривали простые случаи для генераторов и оператор yield, однако, как и все вещи Python, он на этом не заканчивается, вокруг него есть еще кое-что, хотя идея его - это то, что вы до сих пор изучили.

Как мы уже обсуждали, когда мы используем yield, мы сохраняем локальное состояние для функции и возвращаем выражение значения вызывающей функции. Но что мы имеем в виду, чтобы спасти местное государство? ну ... вот тут-то и получается очень интересно. Когда выполняется инструкция Python yield, программа приостанавливает выполнение функции и возвращает полученное значение вызывающей стороне. Когда функция приостанавливается, состояние этой функции сохраняется, включая такие данные, как любые привязки переменных, указатель инструкций, внутренний стек и любую обработку исключений. Когда генератор снова вызывается, состояние восстанавливается, и функция продолжается с последнего оператора yield, который он попал, как если бы предыдущий yield не был вызван и функция не была бы приостановлена.

Довольно аккуратно! Давайте посмотрим на примере, чтобы лучше понять это

>>> def multiple_yield():
...     value = "I'm here for the first time"
...     yield value
...     value = "My Second time here"
...     yield value
...
>>> multi_gen = multiple_yield()
>>> print(next(multi_gen))
I'm here for the first time
>>> print(next(multi_gen))
My Second time here
>>> print(next(multi_gen))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

В первый раз, когда мы выполнили функцию, указатель выполнения был в начале, и, таким образом, мы попали в первый yield в строке 2, таким образом, на экране печатается инструкция «Я здесь впервые». При втором вызове next() указатель выполнения продолжается со строки 3, попадает во второй оператор yield в строке 4 и возвращает «Я здесь второй раз», хотя технически был в этой строке только один раз 😛. Теперь, когда мы вызываем next() в третий раз, мы получаем ошибку. Это связано с тем, что генераторы, как и все итераторы, могут быть исчерпаны, и если вы попытаетесь вызвать next() после этого, вы получите эту ошибку.

Методы опережающего генератора

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

  • .send()
  • .throw()
  • .close()

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

def isPrime(n):
    if n < 2 or n % 1 > 0:
        return False
    elif n == 2 or n == 3:
        return True
    for x in range(2, int(n**0.5) + 1):
        if n % x == 0:
            return False
    return True

def getPrimes():
    value = 0
    while True:
        if isPrime(value):
            yield value
        value += 1

Как использовать .send ()

.send() позволяет установить значение генератора в любое время. Предположим, вы хотите сгенерировать только простые числа, начиная с 1000, и здесь .send() пригодится. Давайте посмотрим на этот пример:

prime_gen = getPrimes()
print(next(prime_gen))
print(prime_gen.send(1000))
print(next(prime_gen))

И когда мы его запускаем, получаем:

2
3
5

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

def getPrimes():
    value = 0
    while True:
        if isPrime(value):
            i = yield value
            if i is not None:
                value = i
        value += 1

Теперь снова бежим

prime_gen = getPrimes()
print(next(prime_gen))
print(prime_gen.send(1000))
print(next(prime_gen))

и получаем:

2
1009
1013

Отлично! Хорошая работа!

Как использовать .throw ()

.throw(), как вы, наверное, догадались, позволяет генерировать исключения с генератором. Это может быть полезно, например, для завершения итерации на определенном значении.

Давайте посмотрим, как это работает:

prime_gen = getPrimes()

for x in prime_gen:
    if x > 10:
        prime_gen.throw(ValueError, "I think it was enough!")
    print(x)

и получаем:

2
3
5
7
Traceback (most recent call last):
  File "test.py", line 25, in <module>
    prime_gen.throw(ValueError, "I think it was enough!")
  File "test.py", line 15, in getPrimes
    i = yield value
ValueError: I think it was enough!

Интересной особенностью этого является то, что ошибка генерируется внутри генератора, что можно увидеть на трассировке стека.

Как использовать .close ()

В предыдущем примере мы останавливаем итерацию, вызывая исключение, однако это не очень элегантно. Лучший способ завершить итерацию - использовать .close().

prime_gen = getPrimes()

for x in prime_gen:
    if x > 10:
        prime_gen.close()
    print(x)

с выходом:

2
3
5
7
11

В этом случае генератор остановился, и мы вышли из цикла, не вызвав исключения.

Заключение

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

Я надеюсь, что теперь вы лучше понимаете генераторы и можете использовать их в своем следующем проекте.

Спасибо за прочтение!

Первоначально опубликовано на https://livecodestream.dev 30 мая 2020 г.