Итераторы и генераторы в Python

Alexander Shepetko5 октября, 15:00 116 0

Об итерируемых объектах, итераторах и генераторах в Python для начинающих

Рано или поздно всякий Python-разработчик сталкивается с осознанием того, что он не до конца понимает смысл и различия следующих понятий:

  • контейнер (container);
  • итерируемый объект (iterable);
  • итератор (iterator);
  • генератор (generator);
  • генераторное выражение (generator expression)
  • списочное включение (list comprehension).

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

Контейнеры

Контейнер — это тип данных, предназначенный для хранения элементов и предоставляющий набор операций для работы с ними. Сами контейнеры и, как правило, их элементы хранятся в памяти. В Python существует масса разнообразных контейнеров, среди которых всем хорошо знакомые:

  • list, deque, …
  • set, frozensets, …
  • dict, defaultdict, OrderedDict, Counter, …
  • tuple, namedtuple, …
  • str

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

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

>>> assert 1 in [1, 2, 3]      # списки
>>> assert 4 not in [1, 2, 3]
>>> assert 1 in {1, 2, 3}      # множества
>>> assert 4 not in {1, 2, 3}
>>> assert 1 in (1, 2, 3)      # кортежи
>>> assert 4 not in (1, 2, 3)

Применительно к словарям операция проверки вхождения работает с ключами словаря:

>>> d = {1: 'foo', 2: 'bar', 3: 'qux'}
>>> assert 1 in d
>>> assert 4 not in d
>>> assert 'foo' not in d  # 'foo' не является _ключом_ словаря

Ну и, наконец, подобным образом вы можете выполнять проверку вхождения подстроки в строку:

>>> s = 'foobar'
>>> assert 'b' in s
>>> assert 'x' not in s
>>> assert 'foo' in s

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

Обратите внимание, что несмотря на то, что большинство контейнеров предоставляют возможность извлекать из них любой элемент, само по себе наличие этой возможности не делает объект контейнером, а лишь итерируемым объектом, о чём будет рассказано дальше. Также, контейнер не обязан быть итерируемым. Например, фильтр Блума предоставляет возможность узнать, содержится ли элемент в структуре данных, но при этом нет возможности извлекать из неё отдельные элементы.

Итерируемые объекты

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

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

>>> x = [1, 2, 3]
>>> y = iter(x)
>>> z = iter(x)
>>> next(y)
1
>>> next(y)
2
>>> next(z)
1
>>> type(x)

>>> type(y)

Здесть х — это итерируемый объект, в то время как y и z два отдельных экземпляра итератора, производящего значения из итерируемого объекта x. Как видим, y и z сохраняют состояние между вызовами next(). В данном примере в качестве источника данных для итератора используется список, но это не является обязательным условием.

Часто, чтобы сократить объем кода, классы итерируемых объектов имплементируют сразу оба метода: __iter()__ и __next()__, при этом __iter()__ возвращает self. Таким образом класс одновременно является и итерируемым и итератором самого себя. Однако, лучшей практикой всё же считается в качестве итератора возвращать отдельный объект.

Итак, когда выполняется следующий код:

x = [1, 2, 3]
for elem in x:
    ...

вот что на самом  деле происходит:

Если диассемблировать код, представленный выше, мы обнаружим вызов GET_ITER, который по сути является следствием вызова iter(x). Инструкция FOR_ITER является эквивалентом многократного вызова next() до тех пор, пока не будет возвращён последний элемент. Этого, правда, не видно в байт-коде из-за оптимизаций, вносимых интерпретатором.

>>> import dis
>>> x = [1, 2, 3]
>>> dis.dis('for _ in x: pass')
  1           0 SETUP_LOOP              14 (to 17)
              3 LOAD_NAME                0 (x)
              6 GET_ITER
        >>    7 FOR_ITER                 6 (to 16)
             10 STORE_NAME               1 (_)
             13 JUMP_ABSOLUTE            7
        >>   16 POP_BLOCK
        >>   17 LOAD_CONST               0 (None)
             20 RETURN_VALUE

Итераторы

Итак, чем же является итератор? Итератор — это вспомогательный объект, возвращающий следующий элемент всякий раз, когда выполняется вызов функции next() с передачей этого объекта в качестве аргумента. Таким образом, любой объект, предоставляющий метод __next()__, является итератором, возвращающим следующий элемент при вызове этого метода, при этом совершенно неважно как именно это происходит.

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

Существует бесчисленное множество примеров использования итераторов. Например, все функции пакета itertools возвращают итераторы. 

Некоторые из которых генерируют бесконечные последовательности:

>>> from itertools import count
>>> counter = count(start=13)
>>> next(counter)
13
>>> next(counter)
14

Некоторые создают бесконечные последовательности из конечных:

>>> from itertools import cycle
>>> colors = cycle(['red', 'white', 'blue'])
>>> next(colors)
'red'
>>> next(colors)
'white'
>>> next(colors)
'blue'
>>> next(colors)
'red'

Или конечные последовательности из бесконечных:

>>> from itertools import islice
>>> colors = cycle(['red', 'white', 'blue'])  # infinite
>>> limited = islice(colors, 0, 4)            # finite
>>> for x in limited:                         # so safe to use for-loop on
...     print(x)
red
white
blue
red

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

>>> class fib:
...     def __init__(self):
...         self.prev = 0
...         self.curr = 1
... 
...     def __iter__(self):
...         return self
... 
...     def __next__(self):
...         value = self.curr
...         self.curr += self.prev
...         self.prev = value
...         return value
...
>>> f = fib()
>>> list(islice(f, 0, 10))
[1, 1, 2, 3, 5, 8, 13, 21, 34, 55]

Обратите внимание, что приведённый выше класс одновременно является итерируемым объектом (поскольку реализует метод __iter__()) и итератором (поскольку реализует метод __next__()).

Состояние этого итератора полностью хранится в переменных prev и curr экземпляра объекта, значение которых используется при каждом последующем обращении к итератору. Каждый вызов next() делает две важные вещи:

  1. модифицирует внутреннее состояние объекта;
  2. возвращает следующий элемент последовательности.

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

Генераторы

Итак, наконец-то мы добрались до самого интересного! Генераторы являются безумно интересной и полезной штукой в Python. Генератор — это особый, более изящный случай итератора.

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

Давайте внесём немного ясности:

  1. любой генератор является итератором (но не наоборот!);
  2. следовательно, любой генератор является "ленивой фабрикой", возвращающей значения последовательности по требованию.

Вот пример итератора последовательности чисел Фибоначчи в исполнении генератора:

>>> def fib():
...     prev, curr = 0, 1
...     while True:
...         yield curr
...         prev, curr = curr, prev + curr
...
>>> f = fib()
>>> list(islice(f, 0, 10))
[1, 1, 2, 3, 5, 8, 13, 21, 34, 55]

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

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

Теперь, когда происходит вызов функции fib

f = fib()

будет создан и возвращён экземпляр генератора.  К данному моменту ещё никакого кода внутри функции не выполняется и генератор ожидает вызова.

Дальше созданный экземпляр генератора передаётся в качестве аргумента функции islice, которая также возвращает итератор, следовательно также никакого кода функции fib пока что не выполняется.

И, наконец, происходит вызов list() с передачей в качестве аргумента итератора, возвращённого функцией islice(). Чтобы list() смогла построить объект списка на основе полученного аргумента, ей необходимо получить все значения из этого аргумента. Для этого list() выполняет последовательные вызовы метода next() итератора, возвращённого вызовом islice(), который, в свою очередь, выполняет последовательные вызовы next() в экземпляре итератора f.

При первом запросе значения итератора f будет выполнен код:

prev, curr = 0, 1

После чего произойдёт вход в тело цикла и оператор yield вернёт первое значение. На этом работа кода внутри f будет приостановлена до следующего вызова next(). Значение, возвращенное при помощи yield, будет передано итератору islice(), который передаст его в list(), таким образом в список добавится значение.

При втором и последующих обращениях к итератору f будет выполняться код с того места, где было прервано выполнение, то есть после вызова оператора yield. После того, как этот когда будет выполнен, выполнение цикла начнётся заново и так же, как и при первом вызове, итератор будет возвращать очередное значение и "засыпать" в аккурат после вызова yield.

Весь описанный процесс повторится 10 раз, до тех пор, пока list() не запросит у islice() одиннадцатый элемент, на что тот ответит выбросом исключения StopIteration, что будет свидетельствовать о том, что достигнут конец последовательности. Обратите внимание, что одиннадцатый элемент не будет сохранён в конечном объекте списка и память, отведённая под него, будет освобождена сборщиком мусора.

Типы генераторов

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

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

>>> numbers = [1, 2, 3, 4, 5, 6]
>>> [x * x for x in numbers]
[1, 4, 9, 16, 25, 36]

Или, то же самое, но в виде множества:

>>> {x * x for x in numbers}
{1, 4, 36, 9, 16, 25}

Или в виде словаря:

>>> {x: x * x for x in numbers}
{1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36}

Или, наконец, используя генераторное выражение (обратите внимание, это НЕ кортеж!):

>>> lazy_squares = (x * x for x in numbers)
>>> lazy_squares
 at 0x10d1f5510>

Заключение

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

def something():
    result = []
    for ... in ...:
        result.append(x)
    return result

И замените их генераторами:

def iter_something():
    for ... in ...:
        yield x