Генераторы Python. Их создание и использование

Генераторы Python

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

Прочитав эту статью, вы узнаете:

  • Что собой представляют генераторы Python и как их использовать
  • Как задавать функции и выражения создающие генераторы
  • Как работает в Python ключевое слово yield

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

По ссылке ниже вы можете скачать копию файла с данными, используемыми в этом руководстве.

Скачать табличные данные, используемые в данном руководстве.

Использование Генераторов

Функции генераторов (их описание можно почитать в PEP 255) представляют собой особый вид функций, которые возвращают «ленивый итератор». И хотя содержимое этих объектов вы можете перебирать также как и списки, но при этом, в отличие от списков, ленивые итераторы не хранят свое содержимое в памяти. Чтобы составить общее представление об итераторах в Python взгляните на статью Python “for” Loops (Definite Iteration).

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

Пример 1: Чтение больших файлов

Списки Python

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

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

for row in csv_gen:
    row_count += 1

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

Глядя на этот пример, вы можете предположить что csv_gen является списком. Для того чтобы заполнить этот список, csv_reader() открывает файл и загружает его содержимое в csv_gen. Затем программа перебирает список, увеличивая значение row_count для каждого следующего ряда.

Это вполне приемлемое решение, но будет ли этот подход работать, если файл окажется слишком большим? А что если файл окажется больше чем вся доступная память, которая есть в нашем распоряжении? Для того чтобы ответить на этот вопрос, давайте предположим, что csv_reder() будет открывать файл и считывать его в массив.

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

Эта функция открывает данный файл и использует file.read() вместе со .split() для того, чтобы добавить каждый ряд данных как отдельный элемент списка. Если бы вы использовали эту версию cvs_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

В этом случае open() возвращает объект генератора, который вы можете «лениво» (не обсчитывая заранее) перебирать ряд за рядом. Тем не менее, file.read().split() загружает все данные в память сразу, вызывая ошибку памяти (MemoryError).

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

Генераторы Python

Давайте взглянем на новое определение функции csv_reader():

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

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

Row count is 64186394

Почему так получилось? Да потому что вы по сути превратили функцию csv_reader() в генератор. Эта версия кода открывает файл, проходит по строкам и извлекает для чтения лишь отдельный ряд, вместо того, чтобы возвращать весь файл целиком.

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

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

Такой способ создания генератора csv_gen является более лаконичным.

Более подробно о yield мы расскажем позже, а пока запомните основные отличия между использованием ключевых слов yield и return:

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

Пример 2: Создание бесконечной последовательности

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

a = range(5)
list(a)

[0, 1, 2, 3, 4]

Создание же бесконечной последовательности стопроцентно потребует от нас использования генератора. Причина проста — ограниченность памяти нашего компьютера.

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

Этот блок кода не велик и хорошо смотрится. Сперва, мы задаем переменную num и создаем бесконечный цикл. Затем мы немедленно извлекаем num с помощью yield в ее исходном состоянии (это во многом повторяет то, что делает range()). После этого мы увеличиваем num на 1.

Если вы попробуете запустить этот код в теле цикла for, то увидите, что на самом деле он бесконечный:

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

0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29
30 31 32 33 34 35 36 37 38 39 40 41 42
[...]
6157818 6157819 6157820 6157821 6157822 6157823 6157824 6157825 6157826 6157827
6157828 6157829 6157830 6157831 6157832 6157833 6157834 6157835 6157836 6157837
6157838 6157839 6157840 6157841 6157842
KeyboardInterrupt
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>

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

Вместо использования Loop, вы также можете использовать на генераторе функцию next(). Это окажется особенно удобным при тестировании работы генератора в консоли:

gen = infinite_sequence()
next(gen)
0
next(gen)
1
next(gen)
2
next(gen)
3

Здесь у нас показан генератор, под названием gen, который мы можем вручную перебирать с помощью вызова функции next(). Это работает как отличная проверка. Она позволяет нам убедиться что генератор выдает результат, который мы от него ожидаем.

Примечание: Когда мы используем next(), Python вызывает метод .__next__(), для функции, которая передается в качестве аргумента. При этом существуют специальные возможности, но разговор о них находится за рамками данной статьи. Если вам интересно, попробуйте поменять аргументы, которые передаются в next() и посмотрите на результат.

Пример 3: Нахождение палиндромов

Вы можете использовать бесконечные последовательности множеством различных способов. Одним из них, который мы отметим особенно, является создание детектора палиндромов. Детектор палиндромов выявляет все последовательности букв и цифр, которые являются палиндромами. Это слова или числа, которые читаются одинаково вперед и назад, как «121» например. Сперва давайте зададим наш числовой детектор палиндромов:

def is_palindrome(num):
    # Skip single-digit inputs
    if num // 10 == 0:
        return False
    temp = num
    reversed_num = 0

    while temp != 0:
        reversed_num = (reversed_num * 10) + (temp % 10)
        temp = temp // 10

    if num == reversed_num:
        return num
    else:
        return False

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

for i in infinite_sequence():
    pal = is_palindrome(i)
    if pal:
        print(pal)

11
22
33
[...]
99799
99899
99999
100001
101101
102201
KeyboardInterrupt
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "<stdin>", line 5, in is_palindrome

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

Примечание: на практике вам вряд ли придется писать свой собственный бесконечный генератор последовательностей, по той простой причине, что есть уже очень эффективный генератор itertools.count() из модуля itertools.

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

Понимание работы генератора Python

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

Функции генераторов выглядят и действуют как обычные функции, но с одной определяющей особенностью. А именно, функция генератора используют ключевое слово yield вместо return. Давайте вспомним функцию генератора, которую мы написали ранее:

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

Это похоже на типичное определение функции, за исключением yield и кода, который следует за ним. Ключевое слово yield применяется там, где значение нужно отправить обратно вызывающей стороне. Но в отличие от return, выхода из функции в данном случае не происходит. Вместо этого, при возврате состояние функции запоминается. Более того, когда next() вызывается для объекта-генератора (явно или неявно в цикле for), ранее полученная переменная num увеличивается, а затем возвращается снова. Поскольку функции генератора похожи на другие функции и действуют подобным образом, вы можете предположить, что выражения создающие генераторы очень похожи на другие выражениях в Python создающие объекты.

Примечание. Если вы хотите больше узнать о генераторах списков, множеств и словарей в Python, можете прочитать статью Эффективное использование генераторов списков (англ).

Создание генератора с помощью выражения

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

Давайте для примера возьмем возведение в квадрат некоторых чисел:

nums_squared_lc = [num**2 for num in range(5)]
nums_squared_gc = (num**2 for num in range(5))

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

nums_squared_lc
[0, 1, 4, 9, 16]
nums_squared_gc
<generator object <genexpr> at 0x107fbbc78>

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

Профилирование эффективности генератора

Ранее мы узнали, что использование генераторов является отличным способом оптимизации памяти. И хотя генератор бесконечной последовательности является наиболее ярким примером этой оптимизации, давайте рассмотрим еще один пример с возведением числа в квадрат и проверим размер полученных объектов.

Вы можете сделать это с помощью вызова функции sys.getsizeof ():

import sys
nums_squared_lc = [i * 2 for i in range(10000)]
sys.getsizeof(nums_squared_lc)
87624
nums_squared_gc = (i ** 2 for i in range(10000))
print(sys.getsizeof(nums_squared_gc))
120

В этом случае размер списка, полученного с помощью выражения составляет 87 624 байта, а размер генератора — только 120. То есть, список занимает памяти в 700 раз больше, чем генератор! Однако нужно помнить одну вещь. Если размер списка меньше доступной памяти на работающей машине, тогда обработка его будет занимать меньше времени, чем аналогичная обработка генератора. Чтобы удостовериться в этом, давайте просуммируем результаты приведенных выше выражений. Вы можете использовать для анализа функцию cProfile.run ():

import cProfile
cProfile.run('sum([i * 2 for i in range(10000)])')
         5 function calls in 0.001 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.001    0.001    0.001    0.001 <string>:1(<listcomp>)
        1    0.000    0.000    0.001    0.001 <string>:1(<module>)
        1    0.000    0.000    0.001    0.001 {built-in method builtins.exec}
        1    0.000    0.000    0.000    0.000 {built-in method builtins.sum}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}


cProfile.run('sum((i * 2 for i in range(10000)))')
         10005 function calls in 0.003 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    10001    0.002    0.000    0.002    0.000 <string>:1(<genexpr>)
        1    0.000    0.000    0.003    0.003 <string>:1(<module>)
        1    0.000    0.000    0.003    0.003 {built-in method builtins.exec}
        1    0.001    0.001    0.003    0.003 {built-in method builtins.sum}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

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

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

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