Итерируемый объект, итератор и генератор

Привет, уважаемые читатели Хабрахабра. В этой статье попробуем разобраться что такое итерируемый объект, итератор и генератор. Рассмотрим как они реализованы и используются. Примеры написан на Python, но итераторы и генераторы, на мой взгляд, фундаментальные понятия, которые были актуальны 20 лет назад и еще более актуальны сейчас, при этом за это время фактически не изменились.

Итераторы

Для начала вспомним, что из себя представляет паттерн «Итератор(Iterator)».
Назначение:

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

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

Существуют два вида итераторов, внешний и внутренний.
Внешний итератор — это классический (pull-based) итератор, когда процессом обхода явно управляет клиент путем вызова метода Next.
Внутренний итератор — это push-based-итератор, которому передается callback функция, и он сам уведомляет клиента о получении следующего элемента.

Классическая диаграмма паттерна “Итератор”, как она описана в небезызвестной книги «банды четырех»:


Aggregate — составной объект, по которому может перемещаться итератор;
Iterator — определяет интерфейс итератора;
ConcreteAggregate — конкретная реализация агрегата;
ConcreteIterator — конкретная реализация итератора для определенного агрегата;
Client — использует объект Aggregate и итератор для его обхода.

Пробуем реализовать на Python классический итератор

Абстрактные классы:

class Aggregate(abc.ABC):

    @abc.abstractmethod
    def iterator(self):
        """
        Возвращает итератор
        """
        pass


class Iterator(abc.ABC):
    def __init__(self, collection, cursor):
        self._collection = collection
        self._cursor = cursor

    @abc.abstractmethod
    def first(self):
        """
        Возвращает итератор к началу агрегата.
        Так же называют reset
        """
        pass

    @abc.abstractmethod
    def next(self):
        """
        Переходит на следующий элемент агрегата.
        Вызывает ошибку StopIteration, если достигнут конец последовательности.
        """
        pass

    @abc.abstractmethod
    def current(self):
        """
        Возвращает текущий элемент
        """
        pass

Конкретная реализация итератора для списка:

class ListIterator(Iterator):
    def __init__(self, collection, cursor):
        """
        :param collection: список
        :param cursor: индекс с которого начнется перебор коллекции.
        так же должна быть проверка -1 >= cursor < len(collection)
        """
        super().__init__(collection, cursor)

    def first(self):
        """
        Начальное значение курсора -1.
        Так как в нашей реализации сначала необходимо вызвать next 
         который сдвинет курсор на 1.
        """
        self._cursor = -1

    def next(self):
        """
        Если курсор указывает на послений элемент, то вызываем StopIteration,
        иначе сдвигаем курсор на 1
        """
        if self._cursor + 1 >= len(self._collection):
            raise StopIteration()
        self._cursor += 1

    def current(self):
        """
        Возвращаяем текущий элемент
        """
        return self._collection[self._cursor]

Конкретная реализация агрегата:

class ListCollection(Aggregate):
    def __init__(self, collection):
        self._collection = list(collection)

    def iterator(self):
        return ListIterator(self._collection, -1)

Теперь мы можем создать объект коллекции и обойти все ее элементы с помощью итератора:

collection = (1, 2, 5, 6, 8)
aggregate = ListCollection(collection)
itr = aggregate.iterator()

# обход коллекции
while True:
    try:
        itr.next()
    except StopIteration:
        break
    print(itr.current())

А так как мы реализовали метод first, который сбрасывает итератор в начальное состояние, то можно воспользоваться этим же итератором еще раз:

# возвращаем итератор в исходное состояние
itr.first()

while True:
    try:
        itr.next()
    except StopIteration:
        break
    print(itr.current())

Реализации могут быть разные, но основная идея в том, что итератор может обходить различные структуры, вектора, деревья, хеш-таблицы и много другое, при этом имея снаружи одинаковый интерфейс.

Протокол итерирования в Python

В книге «банды четырех» о реализации итератора написано:

Минимальный интерфейс класса Iterator состоит из операций First, Next, IsDone и CurrentItem. Но если очень хочется, то этот интерфейс можно упростить, объединив операции Next, IsDone и CurrentItem в одну, которая будет переходить к следующему объекту и возвращать его. Если обход завершен, то эта операция вернет специальное значения(например, 0), обозначающее конец итерации.

Именно так и реализовано в Python, но вместо специального значения, о конце итерации говорит StopIteration. Проще просить прощения, чем разрешения.

Сначала важно определиться с терминами.

Рассмотрим итерируемый объект (Iterable). В стандартной библиотеке он объявлен как абстрактный класс collections.abc.Iterable:

class Iterable(metaclass=ABCMeta):

    __slots__ = ()

    @abstractmethod
    def __iter__(self):
        while False:
            yield None

    @classmethod
    def __subclasshook__(cls, C):
        if cls is Iterable:
            return _check_methods(C, "__iter__")
        return NotImplemented

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

class SomeIterable1(collections.abc.Iterable):
    def __iter__(self):
        pass

class SomeIterable2:
    def __iter__(self):
        pass

print(isinstance(SomeIterable1(), collections.abc.Iterable))
# True
print(isinstance(SomeIterable2(), collections.abc.Iterable))
# True

Но есть один момент, это функция iter(). Именно эту функцией использует например цикл for для получения итератора. Функция iter() в первую очередь для получения итератора из объекта, вызывает его метод __iter__. Если метод не реализован, то она проверяет наличие метода __getitem__ и если он реализован, то на его основе создается итератор. __getitem__ должен принимать индекс с нуля. Если не реализован ни один из этих методов, тогда будет вызвано исключение TypeError.

from string import ascii_letters

class SomeIterable3:
    def __getitem__(self, key):
        return ascii_letters[key]

for item in SomeIterable3():
    print(item)

Итого, итерируемый объект — это любой объект, от которого встроенная функция iter() может получить итератор. Последовательности(abc.Sequence) всегда итерируемые, поскольку они реализуют метод __getitem__

Теперь посмотрим, что с итераторами в Python. Они представлены абстрактным классом collections.abc.Iterator:

class Iterator(Iterable):

    __slots__ = ()

    @abstractmethod
    def __next__(self):
        'Return the next item from the iterator. When exhausted, raise StopIteration'
        raise StopIteration

    def __iter__(self):
        return self

    @classmethod
    def __subclasshook__(cls, C):
        if cls is Iterator:
            return _check_methods(C, '__iter__', '__next__')
        return NotImplemented

__next__ Возвращает следующий доступный элемент и вызывает исключение StopIteration, когда элементов не осталось.
__iter__ Возвращает self. Это позволяет использовать итератор там, где ожидается итерируемых объект, например for.
__subclasshook__ Проверяет наличие у класса метода __iter__ и __next__

Итого, итератор в python — это любой объект, реализующий метод __next__ без аргументов, который должен вернуть следующий элемент или ошибку StopIteration. Также он реализует метод __iter__ и поэтому сам является итерируемым объектом.

Таким образом можно реализовать итерируемый объект на основе списка и его итератор:

class ListIterator(collections.abc.Iterator):
    def __init__(self, collection, cursor):
        self._collection = collection
        self._cursor = cursor

    def __next__(self):
        if self._cursor + 1 >= len(self._collection):
            raise StopIteration()
        self._cursor += 1
        return self._collection[self._cursor]

class ListCollection(collections.abc.Iterable):
    def __init__(self, collection):
        self._collection = collection

    def __iter__(self):
        return ListIterator(self._collection, -1)

Варианты работы:

collection = [1, 2, 5, 6, 8]
aggregate = ListCollection(collection)

for item in aggregate:
    print(item)

print("*" * 50)

itr = iter(aggregate)
while True:
    try:
        print(next(itr))
    except StopIteration:
        break

Функция next() вызывает метод __next__. Ей можно передать второй аргумент который она будет возвращать по окончанию итерации вместо ошибки StopIteration.

itr = iter(aggregate)
while True:
    item = next(itr, None)
    if item is None:
        break
    print(item)

Прежде чем переходить к генераторам, рассмотрим еще одну возможность встроенной функции iter(). Ее можно вызывать с двумя аргументами, что позволит создать из вызываемого объекта(функция или класс с реализованным методом __call__) итератор. Первый аргумент должен быть вызываемым объектом, а второй — неким ограничителем. Вызываемый объект вызывается на каждой итерации и итерирование завершается, когда возбуждается исключение StopIteration или возвращается значения ограничителя.

Например, из функции которая произвольно возвращает 1-6, можно сделать итератор, который будет возвращать значения пока не «выпадет» 6:

from random import randint

def d6():
    return randint(1, 6)

for roll in iter(d6, 6):
    print(roll)

Другие примеры

Небольшой класс ProgrammingLanguages, у которого есть кортеж c языками программирования, конструктор принимает начальное значения индекса по названию языка и функция __call__ которая перебирает кортеж.

class ProgrammingLanguages:

    _name = ("Python", "Golang", "C#", "C", "C++", "Java", "SQL", "JS")

    def __init__(self, first=None):
        self.index = (-1 if first is None else
                      ProgrammingLanguages._name.index(first) - 1)

    def __call__(self):
        self.index += 1
        if self.index < len(ProgrammingLanguages._name):
            return ProgrammingLanguages._name[self.index]
        raise StopIteration

Можем перебрать все языки начиная с C# и до последнего:

for lang in iter(ProgrammingLanguages("C#"), None):
    print(lang)

C первого до C:

pl = ProgrammingLanguages()
for lang in iter(pl, "C"):
    print(lang)

Еще один пример:

# читаем файл до пустой строки
with open('mydata.txt') as fp:
    for line in iter(fp.readline, ''):
        print(line)

Генераторы

С точки зрения реализации, генератор в Python — это языковая конструкция, которую можно реализовать двумя способами: как функция с ключевым словом yield или как генераторное выражение. В результате вызова функции или вычисления выражения, получаем объект-генератор типа types.GeneratorType.

В объекте-генераторе определены методы __next__ и __iter__, то есть реализован протокол итератора, с этой точки зрения, в Python любой генератор является итератором.
Концептуально, итератор — это механизм поэлементного обхода данных, а генератор позволяет отложено создавать результат при итерации. Генератор может создавать результат на основе какого то алгоритма или брать элементы из источника данных(коллекция, файлы, сетевое подключения и пр) и изменять их.

Ярким пример являются функции range и enumerate:

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

Yield

Для начало напишем простой генератор не используя объект-генератор. Это генератор чисел Фибоначчи:

class FibonacciGenerator:
    def __init__(self):
        self.prev = 0
        self.cur = 1

    def __next__(self):
        result = self.prev
        self.prev, self.cur = self.cur, self.prev + self.cur
        return result

    def __iter__(self):
        return self

for i in FibonacciGenerator():
    print(i)
    
    if i > 100:
        break

Но используя ключевое слово yield можно сильно упростить реализацию:

def fibonacci():
    prev, cur = 0, 1
    while True:
        yield prev
        prev, cur = cur, prev + cur

for i in fibonacci():
    print(i)
    if i > 100:
        break

Любая функция в Python, в теле которой встречается ключевое слово yield, называется генераторной функцией — при вызове она возвращает объект-генератор.
Объект-генератор реализует интерфейс итератора, соответственно с этим объектом можно работать, как с любым другим итерируемым объектом.

fib = fibonacci()
print(next(fib))
# 0
print(next(fib))
# 1

for num, fib in enumerate(fibonacci()):
    print('{0}: {1}'.format(num, fib))
    if num > 9:
        break
# 0: 0
# 1: 1
# 2: 1
...

Рассмотрим работу yield:

def gen_fun():
    print('block 1')
    yield 1
    print('block 2')
    yield 2
    print('end')

for i in gen_fun():
    print(i)

# block 1
# 1
# block 2
# 2
# end

1. при вызове функции gen_fun создается объект-генератор
2. for вызывает iter() с этим объектом и получает итератор этого генератора
3. в цикле вызывает функция next() с этим итератором пока не будет получено исключение StopIteration
4. при каждом вызове next выполнение в функции начинается с того места где было завершено в последний раз и продолжается до следующего yield

Происходит приблизительно следующее. Генераторная функция разбивается на части:

def gen_fun_1():
    print('block 1')
    return 1


def gen_fun_2():
    print('block 2')
    return 2


def gen_fun_3():
    print('end')


def gen_fun_end():
    raise StopIteration

Создается стейт-машина в которой при каждом вызове __next__ меняется состояния и в зависимости от него вызывается тот или иной кусок кода. Если в функции yield в цикле, то соответственно состояние стейт-машины зацикливается пока не будет выполнено условие.

Свой вариант range:

def cool_range(start, stop, inc):
    x = start
    while x < stop:
        yield x
        x += inc

for n in cool_range(1, 5, 0.5):
    print(n)
# 1
# 1.5
# ...
# 4.5

print(list(cool_range(0, 2, 0.5)))
# [0, 0.5, 1.0, 1.5]

Генераторное выражение (generator expression)

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

В языках программирования есть такие понятия, как ленивые/отложенные вычисления(lazy evaluation) и жадные вычисления(eager/greedy evaluation). Генераторы можно считать отложенным вычислением, в этом смысле списковое включение(list comprehension) очень похожи на генераторное выражение, но являются разными подходами.

(i for i in range(10000000))

[i for i in range(10000000)]

Первый вариант работает схожим с нашей функцией cool_range образом и может генерировать без проблем любой диапазон. А вот второй вариант создаст сразу целый список, со всеми вытекающими от сюда проблемами.

Yield from

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

Функция похожая на itertools.chain:

def chain(*iterables):
    for it in iterables:
        for i in it:
            yield i

g = chain([1, 2, 3], {'A', 'B', 'C'}, '...')
print(list(g))
# [1, 2, 3, 'A', 'B', 'C', '.', '.', '.']

Но вложенные циклы можно убрать, добавив конструкцию yield from:

def chain(*iterables):
    for it in iterables:
        yield from it

g = chain([1, 2, 3], {'A', 'B', 'C'}, '...')
print(list(g))
# [1, 2, 3, 'A', 'B', 'C', '.', '.', '.']

Основная польза yield from в создании прямого канала между внутренним генератором и клиентом внешнего генератора. Но это уже больше тема про сопрограммы(coroutines), которые заслуживают отдельной статьи. Там же можно обсудить методы генератора: close(), throw() и send().

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

from collections import Iterable

def flatten(items, ignore_types=(str, bytes)):
    """
      str, bytes - являются итерируемыми объектами,
       но их хотим возвращать целыми
    """
    for x in items:
        if isinstance(x, Iterable) and not isinstance(x, ignore_types):
            yield from flatten(x)
        else:
            yield x

items = [1, 2, [3, 4, [5, 6], 7], 8, ('A', {'B', 'C'})]

for x in flatten(items):
    print(x)

# 1
# 2
# 3
# 4
# 5
# 6
# 7
# 8
# A
# C
# B

На сегодня все. Спасибо!