«Напиши, пожалуйста, кастомный итератор,» — такое задание довольно часто дают на собеседованиях, и я раз за разом вижу обреченные глаза кандидата, когда он сталкивается с подобной просьбой. Опыт участия в собеседованиях показал мне, что большинство начинающих разработчиков бегут от этой темы, потому что она кажется слишком запутанной и непонятной. А ведь ничего сложного в ней нет, если подобраться к ней правильным образом — в чём я и постараюсь помочь дорогим читателям.

Наше путешествие мы начнем с того, что вообще такое итератор. Итератор — это некий объект, который в себе реализует интерфейс перебора чего-либо. А говоря рабоче-крестьянским языком — это такая штука, которая в себе описывает правило, по которому мы будем перебирать содержимое той или иной коробки.

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

Конечно же, мы можем хранить наше содержимое тумбочки в любой удобной коллекции, например, в списке:

tumb = ["ножницы", "карандаш", "яблоко", "книга"]

И решать задачу добавления объектов посредством методов списка (через append, например), а задачу перебора — с помощью цикла for:

for obj in tumb:
		print(obj)

«Ну и при чём тут какой-то там итератор?» — спросите меня вы. А что, если я скажу вам, что цикл for работает не совсем так, как вы думаете?

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

На самом деле цикл for взаимодействует не с самим целевым объектом перебора, а с его итератором! В нашем случае — с итератором списка. То есть он как бы говорит: «Эй, объект! Я хочу тебя перебрать, поэтому дай мне то, что описывает правило твоего перебора!» Если объект сможет ответить на это «Вот, держи!» и вернет циклу for некий объект, то объект называется итерируемым и его можно перебирать. Если же он отвечает что-то вроде «Я не понимаю, о чём ты», то программа выдаст ошибку, и это будет означать, что мы попытались перебрать объект, который для этого не предназначен.

Весь фокус в том, что итерируемый объект в случае успеха в качестве ответа на запрос «Дай правило итерации!» возвращает объект итератор.

Как цикл for получает объект-итератор от целевого итерируемого объекта? С помощью неявного вызова встроенной функции iter, в которую в качестве аргумента он передаёт как раз этот самый целевой итерируемый объект. И если в результате будет получен некий объект, то дальнейшая работа будет производиться уже с ним. Давайте посмотрим, что мы получим в качестве результата, если передадим наш список в метод iter:

>>> print(iter(tumb))
<list_iterator object at 0x10d4c53d0>

Видите? Мы получили объект типа list_iterator, инкапсулирующий в себе то самое правило перебора, которое сейчас будет применяться.

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

Как цикл for нажимет эту воображаемую кнопку у итератора? С помощью ещё одной встроенной функцииnext, аргументом которой является объект-итератор, полученный на предыдущем шаге. Результатом этого в штатном случае будет получение очередного значения. Такая процедура будет повторяться многократно до тех пор, пока цикл for не получит сообщение о том, что все значения уже закончились. Что это за сообщение? Это raise ошибки StopIteration.

А теперь давайте взглянем на аналог цикла for, написанный через while:

tumb = ["ножницы", "карандаш", "яблоко", "книга"]

# получаем итератор для итерируемого объекта
it = iter(tumb)

try:
    while True:
        next_val = next(it)
        print("Очередное значение:", next_val)
except StopIteration:
    # явно напечатаем сообщение об окончании итерации,
    # хотя цикл for этого не делает и ошибка просто подавляется
    print("Итерация закончена")
print("Программа завершена")

То есть ответственность за перебор лежит не на цикле for (он просто запрашивает итератор и жмëт в нëм кнопку) и не на самом итерируемом объекте (он лишь должен отдавать свой объект-итератор по запросу), а на итераторе!

Хорошо, с этим понятно, но что там с каким-то кастомным итератором? Это про что вообще история? А эта история про ситуации, когда вам необходимо, чтобы объекты ваших самописных классов тоже можно было итерировать. Давайте представим, что нам пришлось написать для тумбочки отдельный класс, чтобы навесить там много разных дополнительных методов, суть которых нам не важна в описываемом контексте, но важно то, что мы хотели бы иметь возможность перебирать нашу тумбочку, как если бы это был простой список.

Пускай у нашей тумбочки предметы хранятся в составе списковых полей, каждое из которых будет соответствовать тому или иному ящику. У тумбочки будут методы добавления и удаления элементов в ящики. Наш класс будет выглядеть как-то так:

class Tumbochka:
		"""Волшебная тумбочка с тремя ящиками для чего угодно"""

    def __init__(self):
        self.boxes = {
            1: [],
            2: [],
            3: []
        }

    def add_to_box(self, obj, box_num):
        if box_num not in {1, 2, 3}:
            print("Вы ввели неправильный номер ящика!")
        else:
            self.boxes[box_num].append(obj)

    def remove_from_box(self, box_num):
        if box_num not in {1, 2, 3}:
            print("Вы ввели неправильный номер ящика!")
        else:
            return self.boxes[box_num].pop()

    def __str__(self):
        boxes_items = self.boxes[1] + self.boxes[2] + self.boxes[3]
        return ", ".join(boxes_items)

Создадим тумбочку, нагрузим её предметами и выведем информацию на экран:

tumb = Tumbochka()
tumb.add_to_box("ножницы", 1)
tumb.add_to_box("карандаш", 2)
tumb.add_to_box("яблоко", 3)
tumb.add_to_box("книга", 1)
print(tumb)

А теперь вопрос: как нам сделать так, чтобы нашу тумбочку можно было итерировать? Можно, конечно, взять и сделать что-то вроде нового списка, который будет хранить в себе сумму элементов трех ящиков: tumb.boxes[1] + tumb.boxes[2] + tumb.boxes[3] и итерировать его, но это не очень хороший подход. Почему? Давайте представим себе, что у нас есть список с несколькими итерируемыми объектами: списком, множеством и строкой.

my_shiny_list = [
    ["Это", "список", "внутри", "списка"],
    {"Это", "множество", "внутри", "списка"},
    "Это строка внутри списка",
]

И у нас вознкает необходимость добавить к этим товарищам ещё и нашу тумбочку (а может даже и не одну):

my_shiny_list = [
    ["Это", "список", "внутри", "списка"],
    {"Это", "множество", "внутри", "списка"},
    "Это строка внутри списка",
    tumb,
]

Согласитесь, что писать сумму списков внутри вложенного списка — не самое изящное решение, да и сам класс в процессе развития нашего проекта может измениться: например, уберутся или добавятся новые ящики. Придется помнить про все подобные места и вносить правки в них.

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

for some_collection in my_shiny_list:
    for el in some_collection:
        print(el)

В какой-то момент мы получим ошибку TypeError: 'Tumbochka' object is not iterable. Python говорит нам, что мы попытались проитерировать объект, который на приказ «Дай мне свой итератор!» отвечает что-то вроде «Я не понимаю, о чëм ты!».

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

Как это устранить? Нужно сделать так, чтобы встроенная функция iter получала от нашего объекта тумбочки её итератор. Для этого нам потребуется дописать в классе тумбочки магический метод __iter__, назначение которого как раз и состоит в том, чтобы создавать и возвращать в результате своей работы некий объект-итератор.

Но давайте посмотрим повнимательнее — для решения нашей задачи достаточно, чтобы в качестве итератора тумбочки выступал итератор списка суммы трëх ящиков. Для его получения нам нужно будет просто передать эту сумму во встроенную функцию iter и уже результат по работы вернуть в качестве результата магического метода __iter__! Вот как это будет выглядеть:

def __iter__(self):
    # получаем сумму предметов всех ящиков
    boxes_items = self.boxes[1] + self.boxes[2] + self.boxes[3]
    # получаем итератор от списка и возвращаем его
    it = iter(boxes_items)
    return it

Теперь наша тумбочка без проблем сможет быть перебранной наравне с известными встроенными коллекциями и по праву будет являться итерируемым объектом.

И что, это и есть кастомный итератор? Нет! В вышеупомянутом примере мы воспользовались итератором списка в качестве итератора нашей тумбочки. То есть сам итератор не имеет ни малейшего понятия, что его вернули как результат работы какой-то там тумбочки; его задача состоит лишь в том, чтобы перебирать.

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

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

def __iter__(self):
    # получаем сумму предметов всех ящиков
    boxes_items = self.boxes[1] + self.boxes[2] + self.boxes[3]
    # возвращаем очередное значение
    # (пару "объект в ящике тумбочки + адрес в памяти") с помощью yield
    for el in boxes_items:
        yield el, id(el)

Результатом работы магического метода __iter__ будет generator_object, из которого по запросу мы сможем получать очередные пары значений.

Ну, а теперь это кастомный итератор? Всë ещë нет! Необходимость написания кастомного итератора возникает тогда, когда мы хотим тонко управлять процессом перебора. Например, иметь возможность при каком-то событии начать итерацию с самого начала или установить значение указателя перебора на определённый элемент. То есть не перебирать всё подряд, как в цикле for, а управлять вручную нашим процессом через непосредственное взаимодействие с итератором.

Как решить эту задачу? Объект-генератор нам не подходит, использовать итераторы коллекций — тоже. Придëтся писать что-то своё :) Это самое «что-то своё» и будет называться кастомным итератором. Это отдельный класс, объект которого будет возвращаться в качестве результата работы метода __iter__. Давайте напишем простой класс:

class TumbochkaIterator:
    pass

А в классе тумбочки поменяем магический метод __iter__ на вот такую реализацию:

def __iter__(self):
    return TumbochkaIterator()

И попробуем получить итератор тумбочки через встроенную функцию iter:

iter(tumb)

Мы увидим ошибку TypeError: iter() returned non-iterator of type 'TumbochkaIterator', которая любезно сообщает о том, что тот объект, который мы вернули, итератором на самом деле не является.

А что является, спросите вы? А является итератором то, что обладает специальным магическим методом, способным возвращать очередное значение. Именно возвращать очередное значение! Таким магическим методом является метод __next__. Этот метод будет отрабатывать каждый раз, когда объект итератора будет передаваться во встроенную функцию next.

Давайте добавим пустой метод __next__ в наш класс-итератор:

class TumbochkaIterator:
    def __next__(self):
        pass

Теперь посмотрим на результат, который выведет этот код:

>>> print(iter(tumb))
<main.TumbochkaIterator object at 0x10acecee0>

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

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

it = iter(tumb)
print(next(it))
print(next(it))
print(next(it))
print(next(it))

None
None
None
None

Мы получили четыре объектаNone, потому что именно их возвращает нам метод __next__. Давайте сделаем так, чтобы __next__ возвращал единичку:

class TumbochkaIterator:
    def __next__(self):
        return 1

Теперь при попытке запустить код выше мы получим четыре единички.

1
1
1
1

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

class TumbochkaIterator:
    def __init__(self, some_objects):
				self.some_objects = some_objects
				self.current = 0

    def __next__(self):
        if self.current < len(self.some_objects):
            result = self.some_objects[self.current]
            self.current += 1
            return result

class Tumbochka:
    """Волшебная тумбочка с тремя ящиками для чего угодно"""

    def __init__(self):
        self.boxes = {
            1: [],
            2: [],
            3: []
        }

    def add_to_box(self, obj, box_num):
        if box_num not in {1, 2, 3}:
            print("Вы ввели неправильный номер ящика!")
        else:
            self.boxes[box_num].append(obj)

    def remove_from_box(self, box_num):
        if box_num not in {1, 2, 3}:
            print("Вы ввели неправильный номер ящика!")
        else:
            return self.boxes[box_num].pop()

    def __str__(self):
        boxes_items = self.boxes[1] + self.boxes[2] + self.boxes[3]
        return ", ".join(boxes_items)

    def __iter__(self):
        return TumbochkaIterator(self.boxes[1] + self.boxes[2] + self.boxes[3])

А теперь давайте запустим вот этот код:

tumb = Tumbochka()
tumb.add_to_box("ножницы", 1)
tumb.add_to_box("карандаш", 2)
tumb.add_to_box("яблоко", 3)
tumb.add_to_box("книга", 1)

it = iter(tumb)
print(next(it))
print(next(it))
print(next(it))
print(next(it))
print(next(it))
print(next(it))

Вот, что мы получим в результате:

ножницы
книга
карандаш
яблоко
None
None

Что это за None? Откуда это всë взялось? Дело в том, что встроенной функции next не важно, что вернулось в качестве результата работы __next__. Пусть даже это будет None.

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

def to_start(self):
    self.current = 0

def to_current(self, val):
    if val >= len(self.some_objects) or val < 0:
        print("Неверное значение для курсора!")
    else:
        self.current = val

Теперь мы можем при ручной итерации через next обнулять нашу итерацию или перемещать курсор к какому-то другому значению. Например, мы можем сказать, что если тот элемент, который я достал из тумбочки — это ножницы, то следующий элемент я буду пропускать и перемещать курсор на шаг вперед. Можете попробовать написать такое условие в качестве домашнего задания к этой статье ;)

А сейчас я предлагаю запустить процесс итерации нашей тумбочки в цикле for и посмотреть на результат:

for el in tumb:
    print(el)

Давайте взглянем на результат (скорее всего, у вас будет просто None, потому что значения будут лететь очень быстро, я рекомендую добавить sleep(.1) перед вызовом print):

ножницы
книга
карандаш
яблоко
None
None
None
None
None
None

Программа ушла в бесконечный цикл и остановить её можно только с помощью ручного останова. Почему так произошло? Да потому что циклу for тоже не важно, что вернулось в качестве очередного значения из __next__ — он будет жать кнопку «Дай!» до тех пор, пока не возникнет исключение StopIteration. Возникновение этого исключения мы и должны теперь предусмотреть в методе __next__. Давайте сделаем это:

def __next__(self):
    if self.current < len(self.some_objects):
        result = self.some_objects[self.current]
        self.current += 1
        return result
    raise StopIteration 

Перезапустим наш код и взглянем на результат:

ножницы
книга
карандаш
яблоко

Теперь всё работает замечательно! Остался последний штрих: по соглашению объекты-итераторы также должны являться итерируемыми объектами. Принято в качестве итераторов для самих итераторов использовать их самих. Звучит зубодробительно, но я думаю, что взглянув на метод __iter__ в составе нашего итератора вопросы уйдут:

class TumbochkaIterator:
		def __init__(self, some_objects):
				self.some_objects = some_objects
        self.current = 0

    def to_start(self):
        self.current = 0

    def to_current(self, val):
        if val >= len(self.some_objects) or val < 0:
            print("Неверное значение для курсора!")
        else:
            self.current = val

    def __iter__(self):
        return self

    def __next__(self):
        if self.current < len(self.some_objects):
            result = self.some_objects[self.current]
            self.current += 1
            return result
        raise StopIteration

То есть объект возвращает сам себя в качестве итератора, если мы хотим проитерировать сам итератор. Это сделано для того, чтобы у мы могли перебирать объекты кастомных итераторов в цикле for точно так же, как это можно делать с generator_objects или итераторами коллекций, вроде list_iterator.

Подытожим:

• Итерируемый объект — это объект, который можно перебирать.

• За правило перебора отвечает итератор, а не сам объект.

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

• Метод, который возвращает итератор, называется __iter__.

• Объект-итератор должен иметь метод __next__, который возвращает очередное значение.

• Цикл for будет вызывать функцию next от итератора до тех пор, пока не получит исключение StopIteration.

• Возникновение StopIteration — это ответственность итератора, а именно его метода __next__.

• Если StopIteration не возникнет никогда, то мы получим бесконечный цикл.

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

Благодарю за внимание.