Просто о сложном — генераторы в Python

В прошлой статье мы детально разбирались с итераторами и итерируемыми объектами. Тема итераторов не может считаться полностью освещенной, если не поговорить про генераторы. В этой статье мы узнаем, что такое генераторы, когда и для чего их можно и нужно использовать, напишем несколько своих генераторов и создадим свой собственный data pipeline (конвейер данных) с использованием нескольких генераторов.

Что такое генератор? Генератор — это функция, которая возвращает так называемый ленивый итератор. В теории программирования ленивыми вычислениями называются вычисления, которые откладывают исполнение того или иного выражения до того момента, как его значение будет непосредственно востребовано. Ленивый итератор — это объект, который как следует из его названия может быть проитерирован как список, но в отличие от списка, этот объект не хранит свои данные в памяти.

Учитывая сказанное выше про то, что генераторы не хранят данные в памяти, мы уже можем сделать вывод, что одна из наиболее распространенных задач, где применяются генераторы — работам с потоками данных и чтение данных из больших файлов, размер которых, возможно, даже превышает допустимый размер оперативной памяти компьютера. Давайте рассмотрим пример считывания данных файла для последующего вывода в консоль. Предположим, у нас есть файл “test.txt”, состоящий из пяти строк, в каждой из которых записано число от 1 до 5. Вариант без использования генераторов:

>>> def read_file(filename):
>>>     with open(filename, 'r') as fp:
>>>         contents = fp.read()
>>>     return contents
>>> print(read_file('test.txt'))

Данный код выведет в консоль:

1
2
3
4
5

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

>>> def read_file2(filename):
>>>     with open(filename, 'r') as fp:
>>>         for row in fp:
>>>             yield row

>>> rf_generator = read_file2('test.txt')
>>> for row in rf_generator:
>>>     print(row)

Наша функция read_file2 возвращает генератор. Это легко проверить написав:

>>>print(type(rf_generator))

Далее мы работаем с генератором точно так же как мы работаем с итератором, мы проходимся по его содержимому до тех пор, пока генератор себя не исчерпает. Цикл for мы можем заменить на ручной проход по генератору используя тот же метод, что использовали для итераторов: next().

>>> rf_generator = read_file2('test.txt')
>>> print(next(rf_generator))
1

Как видим, вывелось содержимое первой строки нашего файла. Несложно догадаться, что произойдет, если мы вызовем метод next еще несколько раз:

>>> rf_generator = read_file2('test.txt')
>>> print(next(rf_generator))
>>> print(next(rf_generator))
>>> print(next(rf_generator))
>>> print(next(rf_generator))
>>> print(next(rf_generator))
>>> print(next(rf_generator))
1
2
3
4
5
Traceback (most recent call last):
  File "D:/PYTHON/LeetCode/generators.py", line 49, in 
    print(next(rf_generator))
StopIteration

Вывелось все содержимое нашего файла. По исчерпанию генератора, поднято исключение StopIteration.

Важно понимать как работает конструкция yield в функции, которая и создает генератор для нас. В отличие от привычной конструкции return, yield не завершает работу функции, а приостанавливает её работу, запоминая свое текущее состояние: все привязки переменных, позиции указателей, исключения и т.п. После приостановления yield возвращает указанное после yield значение в основную программу. После того как основная программа запросит следующий элемент генератора — будь то цикл или ручной вызов метода next(), функция выйдет из состояния останова и продолжит выполнение ровно со строки следующей за выражением yield. Если в функции после yield будет использован оператор return, то для генератора это будет означать действие: “вызвать исключение StopIteration”.

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

>>> def simple_example():
>>>     str = "line 1"
>>>     yield str
>>>     str2 = "line 2"
>>>     yield str2

>>> gt = simple_example()
>>> print(next(gt))
>>> print(next(gt))
line 1
line 2

При вызове генератора используя метод next функция дойдет до первого yield и вернет нам строку “line 1”, которая и будет выведена в консоль. На этом выполнение функции остановится, но её состояние, текущее положение — сохранится. При повторном вызове метода next с нашим генератором, функция продолжит свое выполнение начиная со следующей строки после первого yield. Присвоит новое значение переменной str2 и вновь остановит свое выполнение на втором yield, верну в программу строку “line 2”, которая и будет выведена в консоль. При последующем запросе следующего элемента генератор через next(gt) мы получим исключение StopIteration — генератор исчерпан.

Мы можем создавать генераторы не только через функции с ключевым словом yield, но и через специальные выражения-генераторы, которые не следует путать с генераторами списков. Генераторы списка позволяют сформировать некоторый список используя следующую конструкцию:

>>> my_list = [elem for elem in range(10) if elem % 2 == 0]
>>> print(my_list)
[2, 4, 6, 8]

Здесь мы создали список всех четных чисел в диапазоне от 1 до 9 включительно. Так реализуются генераторы списков. Если мы хотим реализовать генератор возвращающий четные числа в диапазоне от 1 до 9, то мы запишем конструкцию в следующем виде:

>>> my_gen = (elem for elem in range(1,10) if elem % 2 == 0)
>>> for elem in my_gen:
>>>     print(elem)

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

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

>>> my_list = [elem for elem in range(1, 100000) if elem % 2 == 0]
>>> print(sys.getsizeof(my_list))
203244

Мы увеличили количество обрабатываемых чисел до 100 000. Размер списка, который возвращает наш генератор списка = 203 244 байт. Теперь посмотрим на размер нашего генератора.

>>> my_gen = (elem for elem in range(1, 100000) if elem % 2 == 0)
>>> print(sys.getsizeof(my_gen))
56

56 байт. Разница на лицо, не правда ли?

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

Ранее в статье мы разобрались как с генераторами работает метод next(). Но он не единственный, который можно использовать, работая с генераторами. Оказывается, yield не только может передавать значение того, что идет после этого ключевого слова в основную программу, но с помощью yield можно получать обратно то или иное значение, переданное генератору. Удивительно, но факт. Давайте рассмотрим пример генератора, который получает число и начинает перебирать числа от 1 до полученного числа, возвращая каждую итерацию квадрат числа, при этом инкрементирование будет происходить на случайное число от 1 до 5, переданное в генератор методом send().

>>> def power_master(number):
>>>     start = 1
>>>     while start < number:
>>>         print('yielding = {}'.format(start ** 2))
>>>         inc = (yield start ** 2)
>>>         print('inc = {}'.format(inc))
>>>         start += inc

>>> pw_generator = power_master(100)
>>> print('Current item = {}'.format(next(pw_generator)))
>>> print('Current item = {}'.format(pw_generator.send(random.randint(1, 5))))
>>> print('Current item = {}'.format(pw_generator.send(random.randint(1, 5))))
yielding = 1
Current item = 1
inc = 4
yielding = 25
Current item = 25
inc = 1
yielding = 36
Current item = 36

Я не случайно добавил вывод отладочной информации, чтобы было понятнее, что происходит. Итак, наша функция генератор все так же возвращает квадрат текущего числа, но присваивает результат yield-а в переменную inc, на которую в последующем и будет увеличено текущее число start. В основном коде мы создаем генератор для числа 100 — максимальное число после которого будет вызвано исключение StopIteration. Далее мы обращаемся к генератору и выводим первое число с помощью метода next(). После этого мы уже обращаемся к генератору вызывая метод send, в котором передаем случайное число от 1 до 5. Обратите внимание, что send() не просто передает значение, а выводит функцию из точки останова по аналогии с next() — о чем свидетельствует отладочная информация.

Кроме метода stop() мы можем использовать метод throw():

>>> pw_generator.throw(ValueError("Some value error"))

И метод close():

>>> pw_generator.close()

close() используется для остановки генератора. Будет поднято исключение StopIteration.

В заключении, давайте рассмотрим небольшой пример, реализующий конвейер данных, который читает CSV файл зарплатных ведомостей из трех полей (id, name, salary) и возвращающий сумму выданных зарплат у работников с четными идентификационными номерами. Конечно, всякие операции с чтением больших данных из CSV лучше делать с помощью библиотеки Pandas, но для практики мы сделаем это с помощью генераторов.

Вот такой файл мы имеем на входе — salary.csv:

id,name,salary
1,John,100
2,Mark,500
3,Jim,150
4,Mike,700
5,Jane,500
6,Ann,1000

Наш основной код:

>>> rows_gen = (row for row in open('salary.csv'))
>>> rows_list = (row.rstrip().split(',') for row in rows_gen)
>>> headers = next(rows_list)
>>> workers_dict = (dict(zip(headers, row_list)) for row_list in rows_list)
>>> even_salary = (
>>>     int(worker['salary'])
>>>     for worker in workers_dict
>>>     if int(worker['id']) % 2 == 0
>>> )
>>> print('Sum = {}'.format(sum(even_salary)))
Sum = 2200

Безусловно, в данном коде присутствует избыточное количество итераторов, можно было обойтись и меньшим, но для нашего примера с конвейерами данных я считаю, что чем больше генераторов — тем лучше. Разберем наш код. Первым делом мы определяем генератор, который считывает по одной линии из нашего файла и возвращает линию в основную программу. Далее мы создаем второй генератор rows_list, которые получает считанную первым генератором строку из файла и разбивает её по запятым, возвращая каждую итерацию новый список с данными из текущей строки. Первой строкой у нас идут названия колонок (столбцов), поэтому мы считываем их в список headers используя метод next генератора rows_list. Третьим этапом мы создаем генератор workers_dict, который получает список из генератора rows_list и делает из него словарь с ключами, считанными в headers на прошлом шаге. Наконец, мы формируем финальный итератор even_salary, который получает словарь из генератора workers_dict и возвращает зарплату работника предварительного убедившись, что у работника четный идентификационный номер, и проведя конвертацию из строкового объекта в целочисленный тип. Последней строчкой мы передаем наш генератор even_salary в функцию sum, которая извлекает каждый его элемент, суммирует и возвращает итоговое число.

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