Ітератори в Python

Ітератори — об'єкти, які можна ітерувати один за одним. В Python майже все є ітераторами і розуміння їх роботи допоможе вам писати кращий код. Якщо ви використовували цикл for, ви вже працювали з ними.

Як вони працюють?

Перш ніж ми почнемо писати свої ітератори, нам потрібно розібратися як вони працюють. Коли ми використовуємо цикл for, як Python розуміє що потрібно взяти?

В цьому йому допомагають дві функції: iter та next. Функція iter отримує ітератор з об'єкту. Він просто викликає магічний метод __iter__ у об'єкта. Тобто, якщо ми хочемо, щоб наш об'єкт можна було ітерувати, нам потрібно реалізувати метод __iter__. Коли ітератор отримано, Python викликає функцію next, яка в свою чергу викликає метод __next__ у ітератора. Давайте розглянемо приклад:

>>> l = [1, 2, 3]

>>> i = iter(l)

>>> type(l)
<class 'list'>

>>> type(i)
<class 'list_iterator'>

>>> next(i)
1

>>> next(i)
2

>>> next(i)
3

>>> next(i)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration


>>>

Розглянемо код детальніше. Ми створили список l з трьома елементами. Потім ми викликали iter(l). Тип l — list, але тип i — list_iterator. Цікаво. Тепер ми викликаємо next(i), що повертає нам значення зі списку, одне за одним, а потім викликає виключення StopIteration.

В даному випадку список — наш ітерований об'єкт, а list_iterator — ітератор. Ось так працює цикл for в Python.

l = [1, 2, 3]

iterator = iter(l)

while True:
    try:
        item = next(iterator)
        print(item)
    except StopIteration:
        break

Ітератор

Ітератор — це об'єкт, що реалізує метод __next__, тобто об'єкт, який можна передати в функцію next. Давайте напишемо ітератор, що повертатиме нам цілі числа, до безкінечності.

class InfiniteIterator:
    def __init__(self):
        self.__int = 0

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

Якщо ми будемо викликати next, то отримуватимемо цілі числа, починаючи з одиниці.

>>> inf_iter = InfiniteIterator()
>>> next(inf_iter)
1
>>> next(inf_iter)
2
>>> next(inf_iter)
3
>>> next(inf_iter)
4
>>>

Ітерований об'єкт

Ми хочемо створити об'єкт з властивістю ітерування InfiniteNumbers. Він буде робити наш цикл for нескінченним. Що нам для цього потрібно? У нас є InfiniteIterator, і нам всього лише потрібно оголосити метод __iter__, що буде повертати екземпляр InfiniteIterator.

class InfiniteNumbers:
    def __iter__(self):
        return InfiniteIterator()


infinite_numbers = InfiniteNumbers()

for x in infinite_numbers:
    print(x)

    if x > 99:
        break

Якщо прибрати блок if та break, то цикл буде дійсно нескінченним.

Використання StopIteration

Замість того, щоб зупиняти цикл після ста ітерацій, ми можемо реалізувати це в самому ітераторі за допомогою виключення StopIteration:

class HundredIterator:
    def __init__(self):
        self.__int = 0

    def __next__(self):
        if self.__int > 99:
            raise StopIteration

        self.__int += 1
        return self.__int


class InfiniteNumbers:
    def __iter__(self):
        return HundredIterator()


one_hundred = InfiniteNumbers()

for x in one_hundred:
    print(x)

Ітератори також повинні реалізувати __iter__

Ми бачимо, що __next__ прекрасно виконує свою роботу. Але для ітератора також потрібно реалізувати метод __iter__. Навіщо? Ось що про це каже офіційна документація:

Ітератори також повинні мати метод __iter__(), що повертає сам ітератор, тому кожен ітератор можна використовувати напряму там, де потрібен ітерований об'єкт.

Якщо ми спробуємо використати наш ітератор в циклі, то отримаємо помилку:

class HundredIterator:
    def __init__(self):
        self.__int = 0

    def __next__(self):
        if self.__int > 99:
            raise StopIteration

        self.__int += 1
        return self.__int


one_hundred = HundredIterator()

for x in one_hundred:
    print(x)

Ось таку:

Traceback (most recent call last):
  File "iter.py", line 15, in <module>
    for x in one_hundred:
TypeError: 'HundredIterator' object is not iterable

Ця помилка виникає тому що цикл for спочатку викликає функцію iter на переданому об'єкті. А потім викликає next на ітераторі. І проблема в тому, що ми не маємо методу __iter__. Давайте виправимо це:

class HundredIterator:
    def __init__(self):
        self.__int = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.__int > 99:
            raise StopIteration

        self.__int += 1
        return self.__int


one_hundred = HundredIterator()

for x in one_hundred:
    print(x)

Тепер все працює.

Короткий підсумок

Давайте підсумуємо текст вище: що нам потрібно щоб організувати ітерацію:

• Кожен об'єкт, що може бути ітерованим, повинен мати метод __iter__, що повертає ітератор.
• Ітератор повинен мати метод __next__, що повертає наступне значення при кожному виклику. Коли значення закінчаться, буде викинуто виключення StopIteration.
• Ітератор також повинен мати метод __iter__, що повертає сам ітератор.

Навіщо нам об'єкти з властивістю ітерування?

Ви бачили, що ітератори можуть виступати в ролі ітерованих об'єктів, і в вас могло виникнути питання: а навіщо тоді самі ітеровані об'єкти? Що ж, давайте повернемося до нашого коду з HundredIterator. Спробуйте додати ще один цикл. Що сталося? Так, нічого не сталося. Справа в тому, що ітератори зберігають стан, і тому, закінчивши ітерацію один раз, його стан зберігається. В цьому і плюс ітерованих об'єктів: кожен раз повертати новий ітератор, що і роблять вбудовані типи (наприклад, list).

Чому ітератори такі важливі?

Ітератори дозволяють обробляти великі об'єкти більш ефективно. Уявіть файл розміром 1 ГБ. Уявили? А тепер спробуємо завантажити його в пам'ять? Нічого гарного з цього не вийде, натомість ми можемо написати ітератор, що буде читати лише один рядок, обробляти його, а потім читати інший рядок. Це дуже зручно.