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

Итератор, итерации — всё это в голове вызывает ассоциации с массивами и с циклическими операциями. В данной статье мы разберемся с вопросом, что такое итератор, чем он отличается от итерируемого объекта, как работает цикл for “под капотом” и, наконец, как реализовать свой собственный итератор.

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

Итератор — это то, что можно получить из итерируемого объекта. Мы можем сказать, что итератор — это тип, позволяющий реализовать поток данных и предлагающий средства для продвижения по нему. В Python итератором является любой объект, реализующий так называемый протокол итераторов, который должен содержать в себе два метода: __iter()__ и __next()__. Метод __iter()__ должен возвращать итератор. Метод __next()__ — следующий элемент потока данных.

Давайте рассмотрим небольшой пример. Предположим, у нас есть список arr из пяти элементов.

>>> arr = [1,2,3,4,5]
>>> arr
[1, 2, 3, 4, 5]
>>> type(arr)

Как мы знаем, список — это итерируемый объект. Из любого итерируемого объекта можно получить итератор. Давайте “извлечем” итератор из списка arr.

>>> arr_iterator = arr.__iter__()
>>> type(arr_iterator)

Также мы знаем, что для реализации протокола итератора, объект должен реализоваться метод __next__(), попробуем его вызвать.

>>> arr_iterator.__next__()
1

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

>>> arr_iterator.__next__()
2
>>> arr_iterator.__next__()
3
>>> arr_iterator.__next__()
4
>>> arr_iterator.__next__()
5
>>> arr_iterator.__next__()
Traceback (most recent call last):
  File "", line 1, in 
StopIteration

Как мы видим, итератор в лице метода __next__() предоставил нам инструмент для прохода по потоку данных, в нашем случае — по всем элементам списка, а дойдя до конца — вызвал специальное исключение StopIteration, сигнализирующее о том, что итератор исчерпал все доступные ему значения.

Данная конструкция для вывода элементов списка, конечно, очень объемна. Обычно для этой задачи мы используем цикл for.

>>> arr =  [1,2,3,4,5]
>>> for elem in arr:
>>>     print(elem)
1
2
3
4
5

А как for перебирает элементы? На самом деле, “под капотом”, цикл for — это синтаксический сахар для цикла while. А цикл while как раз-таки работает ни с чем иным как с итератором итерируемого объекта. Давайте посмотрим на то, как отрабатывает цикл for без прикрас:

>>> arr_iterator = arr.__iter__()
>>> while True:
>>>     elem = arr_iterator.__next__()
>>>     print(elem)
1
2
3
4
5
Traceback (most recent call last):
  File "", line 2, in 
StopIteration

Как мы видим, white делает ровно то, что мы сделали руками выше. Извлекает итератор, и пока есть данные выводит каждый элемент. По исчерпанию данных поднимается исключение StopIteration, который в нашем случае и заканчивает white, выведя информацию о нем в консоль, потому что в отличие от for-а, while не обрабатывает данное исключение.

При желании, мы можем записать наш код немного красивее, не прибегая к непосредственному упоминанию методов __iter__() и __next__(), а заменив их на методы-обертки iter() и next().

>>> arr_iterator = iter(arr)
>>> while True:
>>>     print(next(arr_iterator))

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

class PowerMaster:
    """Класс реализует итератор, возвращающий квадраты чисел от заданного до 1"""
    def __init__(self, number):
        self.number = number
    def __iter__(self):
        return self
    def __next__(self):
        if self.number > 1:
            rez = self.number ** 2
            self.number -= 1
            return rez
        else:
            raise StopIteration

Попробуем создать объект нашего класса и пройтись по нему циклом for:

>>> pm = PowerMaster(5)
>>> for item in pm:
>>>     print(item)
25
16
9
4

А теперь тоже самое, но используя while и явную работу с итератором.

>>> pm = PowerMaster(5)
>>> iter = iter(pm)
>>> while True:
>>>     print(next(iter))
25
16
9
4

Как мы видим, наш итератор прекрасно себя чувствует и работает как полагается. В нашем случае объект нашего класса и есть итератор сам по себе, поэтому __iter__() возвращает self.

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