Итераторы в Python

Перейдём к практике. Представим, что на собеседовании вас попросили реализовать аналог функции range. Попробуем быстро его написать.

PythonКопировать

class Range:
    def __init__(self, stop_value: int):
        self.current = -1
        self.stop_value = stop_value - 1

    def __iter__(self):
        return RangeIterator(self)

class RangeIterator:
    def __init__(self, container):
        self.container = container

    def __next__(self):
        if self.container.current < self.container.stop_value:
            self.container.current += 1
            return self.container.current
        raise StopIteration

Вы получили первую версию работающего кода. Запустим код и убедимся в этом.

PythonКопировать

_range = Range(5)
for i in _range:
    print(i)

В итоге вы получаете числа от 0 до 4 на отдельных строках — всё хорошо. Тут же прилетает вопрос: «А как это работает? Расскажите подробнее». Многие на этом вопросе отваливаются, потому что не знают, как на него ответить. Чтобы не потерпеть фиаско на собеседовании, разберёмся подробнее, как всё-таки устроены итераторы в Python.

Как устроены итераторы

Начнём с класса range. У него внутри реализован магический метод __iter__. Он обозначает, что объект этого класса итерабельный, то есть с ним можно работать в цикле for. Ещё говорят, что __iter__ отдаёт итерируемый объект.

Чтобы код действительно отдавал новые данные из range, нужно реализовать соответствующую функцию. Она как раз и называется итератор. RangeIterator — итератор для класса range. Любой итератор будет реализовывать магическую функцию __next__, в которой он должен отдавать новые значения для объектов класса range. Если вы дошли до конца множества значений, то появляется исключение StopIteration.

Но можно ли как-то упростить написанный выше код? Да, можно.

PythonКопировать

class Range2:
    def __init__(self, stop_value: int):
        self.current = -1
        self.stop_value = stop_value - 1

    def __iter__(self):
        return self

    def __next__(self):
        if self.current < self.stop_value:
            self.current += 1
            return self.current
        raise StopIteration

В Python есть возможность объявить объекты класса и итерабельными, и итераторами. Это удобно, но с точки зрения принципов проектирования приложения у такого объекта есть две особенности: он является итератором и при этом выполняет какую-то свою логику. В мире Python это допустимо, но в некоторых других языках вас могут понять неправильно. Будьте бдительны!

Цикл for под капотом

Ещё стоит рассмотреть, как работает цикл for под капотом.

PythonКопировать

iterable = Range2(5)
iterator = iterable.__iter__()
while True:
    try:
        value = iterator.__next__()
        print(value)
    except StopIteration:
        break

Или чуть проще.

PythonКопировать

iterable = Range2(5)
iterator = iter(iterable)
while True:
    try:
        value = next(iterator)
        print(value)
    except StopIteration:
        break

Многие встроенные стандартные функции в Python являются итераторами. Ранее вы уже ознакомились с итератором range. Рассмотрим другие варианты, которые могут встретиться при работе с кодом.

Enumerate

Функция создаёт и добавляет порядковый номер элементам передаваемой коллекции. Возвращаемый результат — кортеж из индекса и элемента коллекции.

PythonКопировать

my_list = ["a", "b", "c", "d", "e"]
enum_list = enumerate(my_list)
next(enum_list)
next(enum_list)

# >>> (0, "a")
# >>> (1, "b")

for idx, val in enumerate(my_list):
    print(idx, val)

# output
# >>> 0 "a"
# >>> 1 "b"
# >>> 2 "c"
# >>> 3 "d"
# >>> 4 "e"

Zip

Создаёт кортеж из передаваемых пар значений.

PythonКопировать

values = [1, 2, 3, 4, 5]
keys = ['a', 'b', 'c', 'd', 'e']
matches = zip(keys, values)
next(matches)

# output
# >>> ('a', 1)

Open

Функция открывает файловый дескриптор на чтение построчно.

PythonКопировать

file_data = open("my_file.txt")
next(file_data)

# >>> 'bar\n'

next(file_data)

# >>> 'baz\n'

В Python очень много итераторов. Они откладывают выполнение работы до момента, пока не будет запрошен следующий элемент с помощью next. Блок программы выполняется в так называемом «ленивом» режиме.