10 итераторов, о которых вы могли не знать

Одним из главных достоинств Python является выразительность кода. Не последнюю роль в этом играет возможность удобной работы с коллекциями и последовательностями различного вида: перебор элементов списка по одному, чтение файла по строкам, обработка всех ключей и значений в словаре. Эти и многие другие подобные задачи в Python помогает решить так называемый протокол итераторов (Iterator protocol). Именно этот протокол обеспечивает работу цикла for, устанавливает по каким объектам можно итерироваться, а по каким нет. Как мы увидим далее, сам язык и стандартная библиотека очень широко используют возможности протокола. В этой статье попробуем отыскать не самые известные, но от этого не менее интересные примеры итераторов и итерируемых объектов, которые предлагает Python.

0. В качестве вступления

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

from itertools import count

# Последовательность натуральных чисел
natural_numbers = count(start=1, step=1)

for number in natural_numbers:
    print(number)  # Напечатает 1, 2, 3, и т.д.

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

Как не сложно догадаться по самому названию протокола, этим вторым видом объектов являются итераторы. Каждый итератор неразрывно связан со своим итерируемым объектом. Единственное, что "знает" итератор — это какой элемент в итерируемом объекте будет следующим при переборе. Этот элемент он и возвращает при вызове функции next:

# Итерируемый объект
iterable = ["?", "?", "?"]

# Получаем итератор с помощью функции iter
iterator = iter(iterable)

# Перебираем элементы с помощью итератора
print(next(iterator))  # ?
print(next(iterator))  # ?
print(next(iterator))  # ?

# Если следующего элемента не существует, 
# то итератор генерирует специальное исключение StopIteration
print(next(iterator))  # ?

Как только итератор сгенерировал исключение StopIteration, он становится бесполезен. Иногда говорят, что такой итератор исчерпан. Чтобы заново перебрать элементы итерируемого объекта нужно получить новый итератор вызовом функции iter(iterable).

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

iterable = "строка — это тоже итерируемый объект"
iterator = iter(iterable)

# Получение итератора возвращает один и тот же объект
print(iter(iterator) is iterator)  # True

# Используем итератор в качестве итерируемого объекта
for letter in iterator:
    print(letter)  # Напечатает буквы с, т, р, о, и т.д.

Из-за того, что итератор может быть исчерпан, использование одного и того же итератора несколько раз может привести к неожиданному результату:

# Продолжение предыдущего примера

# Используем итератор второй раз
for letter in iterator:
    print(letter)  # Не напечатает ничего!

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

Многие встроенные функции работают с итерируемыми объектами. Например, функция sum считает сумму всех элементов переданного итерируемого объекта. Встроенные функции map, filter, и max тоже работают с любыми итерируемыми объектами. И это только малая часть, далее мы увидим ещё больше таких примеров.

Итак, выделим ещё раз наиболее существенные детали:

• Итерируемый объект — это любой объект, который можно передать в функцию iter и получить итератор

• Итератор — это любой объект, который можно передать в функцию next и получить следующий элемент или исключение StopIteration

• Каждый итератор также является итерируемым объектом

• Итератор может быть исчерпан

1. Создаём собственный итератор

Если протокол итераторов широко используется в языке и стандартной библиотеке, то должна быть возможность создать свой собственный итератор. И, конечно, такая возможность в Python есть. Как уже было сказано выше, итератор неразрывно связан со своим итерируемым объектом, т.е. с объектом, элементы которого можно перебрать по одному. Чтобы объект нашего класса был итерируемым, и функция iter корректно работала с ним, можно определить магический метод __iter__, который будет возвращать итератор. В свою очередь у итератора должен быть магический метод __next__, тогда его можно будет передать в функцию next, и наши классы будут соответствовать протоколу.

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

class Fibonacci:
    """Наш итерируемый объект."""
    def __init__(self, n: int):
        self.n = n  # Нужное количество чисел Фибоначчи
        
    def __iter__(self):
        """Возвращаем итератор, чтобы соответствовать протоколу."""
        return FibonacciIterator(self.n)

class FibonacciIterator:
    """Итератор для перебора чисел Фибоначчи."""
    def __init__(self, n: int):
        self.n = n  # Нужное количество чисел Фибоначчи
        self.current, self.next = 0, 1
  
    def __next__(self) -> int:
        if self.n <= 0:
            # Итератор исчерпан
            raise StopIteration()
        self.n -= 1
        current = self.current
        self.current, self.next = self.next, self.next + current
        return current
        
    def __iter__(self):
        """
        Чтобы соответствовать протоколу, каждый итератор должен 
        одновременно быть итерируемым объектом.
        """
        return self

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

# Продолжение предыдущего примера

fibo = Fibonacci(10)  # итерируемый объект
iterator = iter(fibo)
print(isinstance(iterator, FibonacciIterator))  # True

print(next(iterator))  # 0
print(next(iterator))  # 1
print(next(iterator))  # 1
print(next(iterator))  # 2
# ... и т.д. 

for number in fibo:
    print(number)  # Напечатает 0, 1, 1, 2, 3, и т.д.

# Сумма первых 10-ти чисел Фибоначчи
print(sum(fibo))  # 88

2. Генераторы

Создание двух классов только для того, чтобы иметь возможность работать с какой-либо последовательностью, скорее всего, не самый практичный способ. И Python предоставляет гораздо более удобную альтернативу. Если в теле обычной функции использовать ключевое слово yield, то такая функция перестаёт просто возвращать результат при вызове, а становится генераторной функцией. Главное различие здесь в том, что генераторная функция описывает, каким образом нужно сгенерировать последовательность. При вызове такой функции она всегда возвращает специальный объект, который называется генератор. Этот объект хранит своё внутреннее состояние, позволяя вычислять элементы последовательности не все сразу, а по мере необходимости. Работа генератора приостанавливается на каждом операторе yield, а очередное вычисленное значение последовательности возвращается вызывающей стороне.

def fibonacci(n: int):
    current, next = 0, 1
    for _ in range(n):
        yield current  # Точка приостановки работы генератора
        current, next = next, current + next

Думаю, не окажется сюрпризом тот факт, что каждый генератор — это итератор:

# Продолжение предыдущего примера
from collections.abc import Iterator

generator = fibonacci(10)
print(isinstance(generator, Iterator))  # True

# Используем итератор
print(set(generator))  # {0, 1, 2, 3, 5, 8, 13, 21, 34}

Помимо ключевого слова yield для создания генераторов можно использовать выражение yield from, что позволяет генерировать элементы из любого другого итерируемого объекта.

from collections.abc import Iterable

def flatten(items: Iterable):
    for item in items:
        if isinstance(item, Iterable):
            yield from flatten(item)
        else:
            yield item
            
items = [[0, 1, 2], 3, 4, [[5], [6, 7]], [[[8]]], 9]
print(list(flatten(items)))  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Разговор о генераторах был бы неполным, если не упомянуть генераторные выражения (Generator expressions). Такие выражения позволяют декларативно задать создаваемый итератор и сразу присвоить его какой-либо переменной. Причём, как и в случае обычных генераторов, значения последовательности вычисляются не сразу, а по мере необходимости.

# Итерируемый объект
iterable = range(10)

# Создаём итератор с помощью генераторного выражения
iterator = (item for item in iterable if item % 2 == 0)

# Оператор in также может работать через протокол итераторов,
# просто перебирая элементы по одному и сравнивая с проверяемым значением
print(8 in iterator)  # True

# Но нужно помнить, что итератор может быть исчерпан
print(8 in iterator)  # False

3. Протокол последовательностей

Да-да, в Python много всяких полезных протоколов. Так что ещё одним способом определения итерируемых объектов является протокол последовательностей (Sequence protocol). Чтобы объекты класса поддерживали этот протокол, класс должен определить метод __getitem__, который сможет работать с целочисленными аргументами:

class SquareNumbers:
    def __init__(self, n: int):
        self.n = n
  
    def __getitem__(self, index: int):
        if index < 0 or index >= self.n:
            raise IndexError()
        return index * index
    
square_numbers = SquareNumbers(100)
iterator = iter(square_numbers)

print(next(iterator))  # 0
print(next(iterator))  # 1
print(next(iterator))  # 4
# ... и т.д. 

print(list(SquareNumbers(5)))  # [0, 1, 4, 9, 16]

Без сомнения, довольно интересный способ создать итератор, но для чего это может пригодиться на практике — вопрос открытый.

4. Функция iter и два аргумента

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

from random import choice

# Функция choice возвращает случайный элемент из списка
iterator = iter(lambda: choice(["?", "?", "?"]), "?")

# Итератор будет работать пока вызов лямбды не вернёт "?"
print("".join(iterator))  # Например, ????

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

5. Перечисления

Перечисление (Enumeration) — это тип данных, который имеет ограниченное количество значений. Как следует из названия, эти значения перечисляются при объявлении типа. Для создания перечислений в стандартной библиотеке Python есть специальный модуль enum:

from enum import Enum, auto

class Color(Enum):
    RED = "красный"
    GREEN = "зелёный"
    BLUE = "синий"
    BLACK = "чёрный"
    WHITE = "белый"

Здесь Color является перечисляемым типом данных, который состоит из пяти значений: Color.RED, Color.GREEN, и т.д. Но самое интересное здесь то, что сам Color — это тоже итерируемый объект!

# Продолжение предыдущего примера
from collections.abc import Iterable

print(isinstance(Color, Iterable))  # True

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

print([color.value for color in Color])
# Напечатает ["красный", "зелёный", "синий", "чёрный", "белый"]

6. Бескрайний мир itertools

Если вдруг вы не знакомы с модулем стандартной библиотеки itertools, то срочно это исправьте! Этот модуль, без сомнения, может изменить вашу жизнь к лучшему (по крайней мере ту её часть, которая непосредственно касается Python). Модуль предоставляет большое количество различных итераторов, которые можно комбинировать друг с другом и таким образом решать множество задач связанных с операциями над последовательностями. Большинство функций в модуле itertools принимают один или несколько итерируемых объектов в качестве аргументов и возвращают итератор, который определённым образом преобразует исходные последовательности в новую последовательность. Разбор всех возможностей итераторов модуля itertools занял бы целую статью, так что предлагаю посмотреть только несколько небольших примеров.

Скорее всего вы слышали про одну популярную задачу на собеседованиях, которая называется FizzBuzz. Один из вариантов этой задачи заключается в том, чтобы вывести на экран числа от 1 до 100. При этом если число делится нацело на 3, то вместо самого числа нужно вывести строку fizz, если же число делится нацело на 5, то нужно вывести строку buzz. Ну а если число делится нацело и на 3, и на 5, то на экран нужно вывести fizzbuzz. Посмотрим, как можно было бы решить эту задачу с помощью модуля itertools.

from itertools import count, cycle, islice

numbers = count(start=1, step=1)
fizzes = cycle(["", "", "fizz"])
buzzes = cycle(["", "", "", "", "buzz"])

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

# Продолжение предыдущего примера

fizzbuzzes = islice(zip(numbers, fizzes, buzzes), 100)

С помощью функции islice можно создать итератор, смысл работы которого похож на взятие среза списка (some_list[:100]), только в отличие от обычных срезов функцию islice можно использовать для любых итерируемых объектов. Ну а итератор, который создаёт встроенная функция zip, на каждой итерации собирает в кортеж элементы нескольких итерируемых объектов, переданных в функцию zip.

# Продолжение предыдущего примера

for number, fizz, buzz in fizzbuzzes:
    # Напечатает 1, 2, fizz, 4, buzz, fizz, ..., 14, fizzbuzz, и т.д.
    print(f"{fizz}{buzz}" or number) 

Другим интересным примером является итератор, который создаёт функция product. Такое название функции обусловлено тем, что итератор возвращает элементы из декартова произведения переданных в функцию итерируемых объектов. Функция product и создаваемый ею итератор позволяют заменить несколько вложенных циклов for одним, что в некоторых случаях может помочь сделать код более читаемым.

from itertools import product

letters = "abcdefgh"
numbers = [1, 2, 3, 4, 5, 6, 7, 8]

for letter in letters:
    for number in numbers:
        # Напечатает a1, a2, a3, ..., b1, b2, и т.д
        print(f"{letter}{number}")

# Заменяем вложенные циклы одним
for letter, number in product(letters, numbers):
    # Напечатает a1, a2, a3, ..., b1, b2, и т.д
    print(f"{letter}{number}")

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

7. Стандартный ввод

Для интерактивного взаимодействия с пользователем в командной строке обычно используется функция input(), которая блокирует выполнение программы, ожидая ввода пользователя с клавиатуры. После нажатия клавиши Enter введённый пользователем текст можно сохранить в переменную:

name = input("Введите ваше имя: ")
# Ожидание ввода ...
print("Ваше имя:", name)

Помимо функции input() для работы со стандартным потоком ввода в модуле sys есть специальный объект stdin. С этим объектом можно работать, как с обыкновенным файлом открытым для чтения. Например, можно прочитать строку:

import sys

name = sys.stdin.readline().strip()
# Ожидание ввода ...
print("Ваше имя:", name)

Интересно, что как и любой открытый файл (вы же знаете, что и по файлу в Python можно итерироваться?), объект stdin тоже является итерируемым. На каждой итерации этот объект возвращает очередную строку из стандартного потока ввода. Чтобы наглядно проиллюстрировать эту возможность, воспользуемся механизмом перенаправления ввода (в ОС Linux) и одной из самых актуальных задач в современном программировании, а именно подсчётом количества строк в файле:

# Файл lc.py
import sys

# Итерируемся по строкам стандартного потока ввода
line_count = sum(1 for line in sys.stdin)

print("Количество строк:", line_count)

Создадим какой-нибудь файл для тестирования нашей программы:

Файл data.txt
aaa
bbb
ccc

Ну и наконец, используем перенаправление ввода и нашу программу по прямому назначению:

$ python lc.py < data.txt
Количество строк: 4

8. Содержимое каталога

Раз уж мы немного затронули тему файлов в предыдущем разделе, то стоит её продолжить и рассмотреть модуль pathlib — очень полезный модуль стандартной библиотеки для работы с файловой системой. Pathlib предоставляет класс Path, который является абстракцией над путём к файлу в операционной системе. С помощью объектов этого класса можно выполнять множество операций над файлами и каталогами: проверять существование, изменять права доступа, читать содержимое и многое-многое другое. В модуле pathlib остановимся только на одном интересующем нас итераторе, который позволяет перебрать содержимое каталога операционной системы.

from pathlib import Path

directory = Path("/home/habr/files")
print(directory.is_dir())  # True

# Получаем итератор по содержимому каталога
directory_iterator = directory.iterdir()

text_files = (
    # Читаем содержимое каждого текстового файла
    path.read_text()  
    for path in directory_iterator
    if (
        # Проверяем, что путь ведёт к текстовому файлу
        path.is_file() and path.suffix == ".txt"
    )
)

# Используем содержимое выбранных текстовых файлов
for text in text_files:
    print(text)

9. CSV

CSV — это популярный текстовый формат для представления табличных данных. Каждая строка в таком файле — это одна строка таблицы. Для работы с файлами этого формата в стандартной библиотеке Python есть отдельный модуль (вы всё ещё удивлены?), который так и называется — csv. Ну а раз есть специальный модуль, есть набор табличных строк в CSV файле, то должен быть и итератор для перебора этих строк. И, конечно, такой итератор в модуле csv есть. Давайте посмотрим следующий пример, файл data.csv выглядит следующим образом:

Имя,Возраст,Любимый цвет
Алиса,22,зелёный
Маргарита,29,белый
Арина,13,жёлтый
Полина,31,фиолетовый

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

import csv
from itertools import islice

columns = ["Имя", "Возраст", "Любимый цвет"]

with open("data.csv", "r") as csv_file:
    csv_reader = csv.DictReader(csv_file, columns)
    
    # Получаем итератор по строкам CSV файла
    csv_iterator = iter(csv_reader)
    
    # Пропускаем первую строку с названиями колонок
    next(csv_iterator)

    # Итерируемся по строкам
    for row in csv_iterator:
        # Напечатает Алиса 22 зелёный, Маргарита 29 белый, и т.д.
        print(row["Имя"], row["Возраст"], row["Любимый цвет"])

10. Пул процессов

Для ускорения сложных вычислительных задач может использоваться одновременное выполнения кода несколькими ядрами процессора. Из-за особенностей интерпретатора Python (пресловутый GIL) параллельные вычисления обычно реализуются через процессы операционной системы. Для работы с процессами в стандартной библиотеке Python существует модуль multiprocessing, в котором нас будет интересовать класс Pool, позволяющий легко создать процессы для параллельных вычислений на нескольких ядрах процессора. Предположим, наша сложная вычислительная задача выглядит следующим образом:

import time

def do(n: int) -> int:
    time.sleep(3)  # Доооолгие вычисления ...
    return n * n

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

# Продолжение предыдущего примера

tasks = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

Как легко посчитать, последовательное вычисление всех задач займёт около 36 секунд:

# Продолжение предыдущего примера

t0 = time.perf_counter()

results = list(map(do, tasks))
print(results)  # [1, 4, 9, ..., 144]

# Напечатает Время выполнения 36 секунд
print("Время выполнения", int(time.perf_counter() - t0), "секунд")

Попробуем ускорить наши вычисления с помощью пула процессов из модуля multiprocessing:

# Продолжение предыдущего примера
from multiprocessing import Pool

t0 = time.perf_counter()

# Создаём пул из четырёх процессов
with Pool(processes=4) as pool:
    results_iterator = pool.imap(do, tasks)
    results = list(results_iterator)
    print(results)  # [1, 4, 9, ..., 144]

# Напечатает Время выполнения 9 секунд
print("Время выполнения", int(time.perf_counter() - t0), "секунд")

Как видно из примера, метод пула процессов imap возвращает итератор по результатам вычислений. При этом всю низкоуровневую работу по управлению процессами и распределению задач между ними берёт на себя сам пул, обеспечивая удобный способ работы с параллельными вычислениями.

А знаете ли вы ещё какие-нибудь интересные примеры итераторов и итерируемых объектов в стандартной библиотеке Python? Пишите в комментариях!