Как создать свой итератор в Python

Некоторое время назад я написал статью о протоколе итераторов, который поддерживает цикл for в Python. Одна вещь, которую я пропустил в этой статье, это то, как создавать свои собственные итераторы.

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

Что такое итератор?

Сначала давайте быстро рассмотрим, что такое итератор. Для более подробного объяснения посмотрите мой доклад Loop Better или прочитайте статью на основе этого доклада.

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

Итератор (iterator) — это объект, который выполняет фактический проход по элементам.

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

>>> favorite_numbers = [6, 57, 4, 7, 68, 95]
>>> iter(favorite_numbers)
<list_iterator object at 0x7fe8e5623160>

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

>>> favorite_numbers = [6, 57, 4, 7, 68, 95]
>>> my_iterator = iter(favorite_numbers)
>>> next(my_iterator)
6
>>> next(my_iterator)
57

Есть еще одно правило, которое делает все более интересным: итераторы сами по себе также являются итерируемыми объектами. Я объяснил последствия этого более подробно в докладе Loop Better, о которой я упоминал выше.

Зачем нужны итераторы?

Итераторы позволяют вам сделать поочередно перебрать элементы, которые будут вычисляется по мере их поступления. Использование итератора вместо списка list, набора set или другой итерируемой структуры данных может иногда позволить нам сэкономить память. Например, мы можем использовать itertools.repeat, чтобы создать итерируемый объект с большим количество элементов:

>>> from itertools import repeat
>>> lots_of_fours = repeat(4, times=100_000_000)

Этот итератор занимает всего 56 байт памяти на моей машине:

>>> import sys
>>> sys.getsizeof(lots_of_fours)
56

Эквивалентный список из 400 миллионов элементов занимает гораздо больше памяти:

>>> lots_of_fours = [4] * 100_000_000
>>> import sys
>>> sys.getsizeof(lots_of_fours)
800000064

Хотя итераторы могут экономить память, они также могут экономить время. Например, если вы хотите распечатать только первую строку 10-гигабайтного файла логов, вы можете легко сделать так:

>>> print(next(open('giant_log_file.txt')))
This is the first line in a giant file

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

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

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

Например, утилита itertools.count предоставит нам итератор, который будет возвращать каждый раз число от 0 и выше:

>>> from itertools import count
>>> for n in count():
...     print(n)
...
0
1
2
(это будет продолжаться бесконечно)

Создание итератора: объектно-ориентированный путь

Итак, мы увидели, что итераторы могут сэкономить нам память, и сэкономить процессорное время.

Давайте создадим наш собственный итератор. А начнем мы с созданием объекта итератора такого же как itertools.count.

Вот итератор Count, реализованный с использованием класса:

class Count:

    """Iterator that counts upward forever."""

    def __init__(self, start=0):
        self.num = start

    def __iter__(self):
        return self

    def __next__(self):
        num = self.num
        self.num += 1
        return num

Этот класс имеет конструктор, который инициализирует наше текущее число в 0 (или что-либо переданное в качестве start). То, что делает этот класс в качестве итератора, это методы __iter__ и __next__.

Когда итерируемый объект передается во встроенную функцию str, вызывается его метод __str__. Когда объект передается во встроенную функции len, вызывается его метод __len__.

>>> numbers = [1, 2, 3]
>>> str(numbers), numbers.__str__()
('[1, 2, 3]', '[1, 2, 3]')
>>> len(numbers), numbers.__len__()
(3, 3)

Встроенной функции iter для итерируемого объекта попытается вызвать его метод __iter__. Встроенная функции next для объекта попытается вызвать его метод __next__.

Предполагается, что функция iter возвращает итератор. Поэтому наша функция __iter__ должна возвращать итератор. Но наш объект сам является итератором, поэтому должен вернуть себя. Поэтому наш объект Count возвращает self из своего метода __iter__.

Функция next должна возвращать следующий элемент в нашем итераторе или вызывать исключение StopIteration, когда элементов больше нет. Мы возвращаем текущий номер и увеличиваем его, чтобы он был больше во время следующего вызова __next__.

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

>>> c = Count()
>>> next(c)
0
>>> next(c)
1

Мы также можем использовать объект Count в цикле for, как с любым другим итератором:

>>> for n in Count():
...     print(n)
...
0
1
2
(это будет продолжаться бесконечно)

Это хороший объектно-ориентированный подход к созданию итератора, но это не обычный способ, которым программисты Python делают итераторы. Обычно, когда нам нужен итератор, мы создаем генератор.

Генераторы: простой способ сделать итератор

Самый простой способ создать свои собственные итераторы в Python — это создать генератор.

Есть два способа сделать генераторы в Python.

Возьмем в качестве источника данных этот список номеров:

>>> favorite_numbers = [6, 57, 4, 7, 68, 95]

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

>>> def square_all(numbers):
...     for n in numbers:
...         yield n**2
...
>>> squares = square_all(favorite_numbers)

Или мы можем сделать такой же генератор, как этот:

>>> squares = (n**2 for n in favorite_numbers)

Первая называется функцией генератора, а вторая называется выражением генератора.

Оба этих объекта-генератора работают одинаково. Оба имеют тип генератора и оба являются итераторами, которые возвращают квадраты чисел из нашего списка чисел.

>>> type(squares)
<class 'generator'>
>>> next(squares)
36
>>> next(squares)
3249

Мы рассмотрим обоих этих подходах к созданию генератора подробнее, но сначала поговорим о терминологии.

Слово «генератор» используется в Python довольно часто:

• Генератор, также называемый объект генератор, является итератором, тип которого является generator.
• Функция генератора — это специальный синтаксис, который позволяет нам создать функцию, которая возвращает объект генератора, когда мы его вызываем.
• Выражение генератора — это синтаксис, похожий на синтаксис генератора списков (comprehension list), который позволяет создавать встроенный объект генератора.

С учетом этой терминологии давайте рассмотрим каждую из этих вещей в отдельности. Сначала рассмотрим функции генератора.

Функции генератора

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

Обычно, когда вы вызываете функцию, ее код сразу выполняется:

>>> def gimme4_please():
...     print("Let me go get that number for you.")
...     return 4
...
>>> num = gimme4_please()
Let me go get that number for you.
>>> num
4

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

>>> def gimme4_later_please():
...     print("Let me go get that number for you.")
...     yield 4
...
>>> get4 = gimme4_later_please()
>>> get4
<generator object gimme4_later_please at 0x7f78b2e7e2b0>
>>> num = next(get4)
Let me go get that number for you.
>>> num
4

Простое присутствие оператора yield превращает функцию в функцию генератора. Это немного странно, но именно так работают функции генератора.

Хорошо, давайте посмотрим на реальный пример функции генератора. Мы сделаем функцию-генератор, которая будет делать то же самое, что и наш класс итератора Count, который мы сделали ранее.

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

Так же, как и наш класс итератора Count, мы можем вручную зациклить генератор, который мы получаем после вызова count:

>>> c = count()
>>> next(c)
0
>>> next(c)
1

И мы можем зациклить этот объект генератора, используя цикл for, как и раньше:

>>> for n in count():
...     print(n)
...
0
1
2
...

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

Выражение генератор

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

Допустим, у нас есть генератор списка, который фильтрует пустые строки из файла и удаляет новые строки с конца:

lines = [
    line.rstrip('\n')
    for line in poem_file
    if line != '\n'
]

Мы могли бы создать генератор вместо списка, превратив квадратные скобки этого выражения в круглые скобки:

ines = (
    line.rstrip('\n')
    for line in poem_file
    if line != '\n'
)

Подобно тому, как наш генератор списков возвращает нам список, наше выражение генератора возвращает нам объект генератора:

>>> type(lines)
<class 'generator'>
>>> next(lines)
' This little bag I hope will prove'
>>> next(lines)
'To be not vainly made--'

Выражения генератора используют более короткий встроенный синтаксис по сравнению с функциями генератора. Хотя они не такие мощные.

Вы можете написать свою функцию генератора в такой форме:

def get_a_generator(some_iterable):
    for item in some_iterable:
        if some_condition(item):
            yield item

Затем вы можете заменить его на выражение генератора:

def get_a_generator(some_iterable):
    return (
        item
        for item in some_iterable
        if some_condition(item)
    )

Выражения генератора против функции генератора

Вы можете думать о выражениях генератора как о списочном понимании мира генератора.

Если вы не знакомы с генераторами списков (list comprehension), я рекомендую прочитать мою статью list comprehensions in Python. В этой статье я отмечаю, что вы можете копировать и вставлять свой код из цикла for в list comprehension.

Вы также можете скопировать и вставить свой код из функции генератора в функцию, которая возвращает выражение генератора:

Выражения генератора относятся к функциям генератора, а генераторы списков (list comprehension) — к простому циклу for с добавлением и условием.

Выражения генератора настолько похожи на выражения, что вы даже можете испытать соблазн сказать «generator comprehension» по аналогии с list comprehension вместо выражения генератора. В русском языке нет устойчивого перевода названия generator comprehension. Хотя технически это не правильное имя, но если вы скажете это, все поймут, о чем вы говорите. Нед Бэтчелдер фактически предложил, чтобы мы все начали называть выражения генератора как generator comprehensions , и я склонен согласиться с тем, потому что это будет более понятное имя.

Так какой же способ лучший что бы создать итератор?

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

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

Я бы порекомендовал использовать выражения генератора так же, как вы используете списки. Если вы выполняете простую операцию мапирования или фильтрации, генераторное выражение — отличное решение. Если вы делаете что-то более сложное, вам, вероятно, понадобится функция генератора.

Я бы порекомендовал использовать функции генератора так же, как вы используете цикл for для создания списка. Везде, где вы видите метод append, вы часто можете использовать вместо этого выражение yield.

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

Генераторы тоже могут помочь при создании итерируемых объектов

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

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

Например, вот итерируемый объект, который предоставляет координаты x-y:

class Point:
    def __init__(self, x, y):
        self.x, self.y = x, y
    def __iter__(self):
        yield self.x
        yield self.y

Обратите внимание, что наш класс Point здесь создает итерацию при вызове (не итератор). Это означает, что наш метод __iter__ должен возвращать итератор. Самый простой способ создать итератор — создать функцию генератора, так что это именно то, что мы и сделали.

Мы вставили yield в наш __iter__, чтобы превратить его в функцию генератора, и теперь наш класс Point может быть зациклен, как и любая другая итерация.

>>> p = Point(1, 2)
>>> x, y = p
>>> print(x, y)
1 2
>>> list(p)
[1, 2]

Функции генератора естественным образом подходят для создания методов __iter__ в ваших итерируемых классах.

Генераторы — это способ создания итераторов.

Словари являются типичным способом составления мапирования в Python. Функции — это типичный способ создания вызываемого объекта в Python. Аналогично, генераторы являются типичным способом создания итератора в Python.

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

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

Оригинал: How to make an iterator in Python