Итераторы в Python
В этой статье вы узнаете, что такое итераторы, зачем они нужны в Python и как создать пользовательский итератор с помощью протокола итератора.
Что такое итератор
Итератор — это объект, у которого есть:
- метод
__iter__
, который возвращает сам объект;- метод
__next__
, который возвращает следующий элемент. Если все элементы уже возвращены, метод вызывает исключение StopIteration.
Методы __iter__ и __next__ называются протоколом итератора.
Python позволяет использовать итераторы в циклах for, списковых включениях и других встроенных функциях, в том числе map
, filter
, reduce
и zip
.
Пример итератора
Ниже — пример класса-итератора Square
, который возвращает квадраты чисел.
class Square:
def __init__(self, length):
self.length = length
self.current = 0
def __iter__(self):
return self
def __next__(self):
if self.current >= self.length:
raise StopIteration
self.current += 1
return self.current ** 2
Как это работает
- Инициализируем атрибуты
length
иcurrent
в методе__init__
. Атрибутlength
определяет количество квадратных чисел, которые должен вернуть класс. А атрибутcurrent
отслеживает текущее целое число. - Создаем метод
__iter__
, который возвращает объектself
. - Создаем метод
__next__
, который возвращает квадрат следующего числа. Если количество возвращенных квадратов чисел большеlength
, метод__next__
вызывает исключение StopIteration.
Используем итератор
Давайте используем только что созданный итератор Square
из примера. Ниже показано, как это сделать в цикле for:
square = Square(5)
for sq in square:
print(sq)
Как это работает
- Создаем новый экземпляр класса
Square
. - Используем цикл for для перебора элементов итератора
Square
.
Как только вы переберете все элементы, итератор будет «исчерпан». Это означает, что вам нужно создать новый итератор для повторного перебора элементов.
Если вы попытаетесь использовать итератор, который уже исчерпан, вы получите исключение StopIteration. Например:
next(square) # ошибка StopIteration
Кроме того, итератор нельзя перезапустить, потому что у него есть только метод __next__
, который возвращает следующий элемент из коллекции.
Что нужно запомнить
- Итератор — это объект, реализующий методы
__iter__
и__next__
. - Итератор нельзя повторно использовать после возврата всех элементов.