Итераторы и генераторы

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

Итераторы

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

Функция iter() принимает итерируемый объект (например, список, кортеж или строку) и возвращает его итератор. Он не является копией переданной коллекции, а просто указывает на то, где в данный момент находится процесс итерации.

Функция next() принимает итератор и возвращает следующий элемент из его последовательности.

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

sweets = iter(["Шоколад", "Сладкая вата"])
print(next(sweets))
# Вывод: Шоколад
print(next(sweets))
# Вывод: Сладкая вата
print(next(sweets))
# Ошибка: StopIteration

Когда элементы в итераторе заканчиваются, функция next() вызывает исключение StopIteration. Это исключение не является ошибкой в обычном смысле; оно просто сигнализирует Python, что итерация завершена. Цикл for автоматически обрабатывает это исключение, чтобы корректно завершить свою работу.

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

sweets = iter(["Шоколад", "Сладкая вата"])
print(next(sweets, "Сладости закончились"))
# Вывод: Шоколад
print(next(sweets, "Сладости закончились"))
# Вывод: Сладкая вата
print(next(sweets, "Сладости закончились"))
# Вывод: Сладости закончились

Вместе функции iter() и next() формируют основу для ручного прохождения по коллекции. С их помощью мы даже можем реализовать цикл for через цикл while:

books = iter(["Метаморфозы", "Государь", "1984"])
while True:
    try:
        book = next(books)
        print(book)
    except StopIteration:
        break
# Вывод: Метаморфозы
# Вывод: Государь
# Вывод: 1984 

Генераторы

Мы уже создавали генераторы с помощью генераторных выражений:

numbers = (n ** 2 for n in range(1, 1000000000))

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

Так как генератор является особым видом итератора, каждый его элемент можно последовательно получить с помощью функции next():

numbers = (n ** 2 for n in range(1, 1000000))
print(next(numbers))
# Вывод: 1
print(next(numbers))
# Вывод: 4
print(next(numbers))
# Вывод: 9

Хотя более эффективно использовать цикл for, так как тогда не требуется вручную обрабатывать исключение StopIteration.

numbers = (n ** 2 for n in range(1, 1000000000))
for n in numbers:
    if n >= 100:
        print(n)
    if n == 144:
        break
# Вывод: 100
# Вывод: 121
# Вывод: 144

Так мы можем получить первые элементы последовательности из миллиарда элементов без необходимости её создания, то есть сэкономив время и память.

Функции-генераторы

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

Оператор yield – это то, что отличает генератор от обычной функции.

Когда Python встречает yield, он делает следующее:

• Приостанавливает выполнение функции.
• Возвращает значение, указанное после yield.
• Запоминает все локальные переменные и текущую позицию в коде.

Например, напишем простую функцию-генератор command_generator():

def command_generator():
    yield "На старт!"
    yield "Внимание!"
    yield "Марш!"

Вызов функции-генератора не возвращает значение напрямую, а создаёт объект-генератор:

commands = command_generator()
print(commands)
# Вывод: <generator object command_generator at 0x...>

Но при каждом вызове генератора, например, в цикле for или с помощью next(), выполнение функции возобновляется с того же места, где оно было приостановлено, до следующего yield:

print(next(commands))
# Вывод: На старт!

print(next(commands))
# Вывод: Внимание!

print(next(commands))
# Вывод: Марш!

Это позволяет создавать последовательности значений по мере необходимости.

Вызов функции-генератора создаёт объект-генератор. Многократный вызов функции-генератора создаёт несколько независимых генераторов. Например, здесь функция всегда будет возвращать "На старт!":

print(next(command_generator()))
# Вывод: На старт!
print(next(command_generator()))
# Вывод: На старт!
print(next(command_generator()))
# Вывод: На старт!

Причина этого заключается в том, что три вызова функции command_generator() создают три совершенно независимых генератора. Поэтому вызов функции next() для каждого возвращает его первое значение.

Генератор для чисел Фибоначчи

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

def fibonacci_generator():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

Здесь функция fibonacci_generator() определена как генератор и создаёт две переменные a и b со значениями 0 и 1 – двумя первыми числами Фибоначчи.

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

Например, выведем на экран первые десять чисел Фибоначчи:

fib_gen = fibonacci_generator()

for _ in range(10):
    print(next(fib_gen), end=" ")
# Вывод: 0 1 1 2 3 5 8 13 21 34

Здесь функция-генератор fibonacci_generator() не создаёт и не хранит весь список чисел, а генерирует их по одному при каждой итерации в цикле for.

В Python подчеркивание (_) может использоваться в качестве имени неиспользуемой переменной для обозначения того, что значение переменной не понадобится в дальнейшем. Но это не правило, а соглашение об именовании, которое не влияет на работу программы, но позволяет не загромождать пространство имён лишними переменными.

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

print(f"\nОдиннадцатое число Фибоначчи: {next(fib_gen)}")
# Вывод: Одиннадцатое число Фибоначчи: 55
print(f"Двенадцатое и тринадцатое числа Фибоначчи: {next(fib_gen)} и {next(fib_gen)}")
# Вывод: Двенадцатое и тринадцатое числа Фибоначчи: 89 и 144

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

Примеры

Пример 1. Управление списком задач

Задачи из списка в системе управления проектами выполняются по одной с явным указанием, когда нужно перейти к следующей.

tasks = [
    "Позвонить Сталину И. В.", 
    "Забрать посылку", 
    "Ответить на письма", 
    "Написать отчёт"
]

# Создаем итератор из списка задач
task_iterator = iter(tasks)

# Выполняем задачи по одной
print("Выполняем следующие задачи:")
print(f"  - {next(task_iterator)}")
print(f"  - {next(task_iterator)}")

# Проверяем, есть ли еще задачи
print(f"Осталось задач: {len(list(task_iterator))}")

Вывод:

Выполняем следующие задачи:
  - Позвонить Сталину И. В.
  - Забрать посылку
Осталось задач: 2

Пример 2. Генератор уникальных идентификаторов

Для новых пользователей система с помощью генератора id_generator() последовательно назначает уникальные идентификаторы. Это позволяет избежать необходимости создавать и хранить список всех возможных ID в памяти:

def id_generator():
    """Генерирует последовательность уникальных ID.
    """
    last_id = 0
    while True:
        last_id += 1
        yield last_id


# Создаем объект-генератор для ID
user_ids = id_generator()

# Получаем новые ID по мере необходимости
print(f"Новый пользователь зарегистрирован с ID: {next(user_ids)}")
print(f"Новый пользователь зарегистрирован с ID: {next(user_ids)}")
print(f"Новый пользователь зарегистрирован с ID: {next(user_ids)}")

Вывод:

Новый пользователь зарегистрирован с ID: 1
Новый пользователь зарегистрирован с ID: 2
Новый пользователь зарегистрирован с ID: 3

Пример 3. Пагинация на сайте

Пагинация на сайте – это метод разбиения большого объема информации на несколько страниц, делая сайт более удобным. Генератор paginate() имитирует вывод результатов поиска и на каждой итерации возвращает срез большого списка данных data, то есть разбивает его на меньшие списки (страницы) по не более чем page_size элементов в каждом:

def paginate(data: list, page_size: int):
    """Генерирует страницы данных из большого списка.
    
    Параметры:
        data: Список данных.
        page_size: Количество элементов на странице.
        
    Yields:
        Список элементов на следующей странице.
    """
    for i in range(0, len(data), page_size):
        yield data[i:i + page_size]


all_results = list(range(1, 46))  # 45 результатов

# Разбиваем результаты на страницы по 10 элементов
pages = paginate(all_results, 10)

# Проходим по страницам и выводим их
for i, page in enumerate(pages):
    print(f"Страница {i+1}: {page}")

Вывод:

Страница 1: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
Страница 2: [11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
Страница 3: [21, 22, 23, 24, 25, 26, 27, 28, 29, 30]
Страница 4: [31, 32, 33, 34, 35, 36, 37, 38, 39, 40]
Страница 5: [41, 42, 43, 44, 45]

Итоги

• Функция iter() преобразует итерируемый объект в итератор.
• Функция next() возвращает следующий элемент итератора.
• Генераторы не хранят всю последовательность в памяти, а генерируют значения по одному.
• Функции-генераторы похожи на обычные функции, однако вместо ключевого return они используют ключевое слово yield.
• Оператор yield приостанавливает выполнение функции, возвращает значение, указанное после yield, и запоминает все локальные переменные и текущую позицию в коде.
• При каждом вызове генератора, например, в цикле for или с помощью next(), выполнение функции возобновляется с того же места, где оно было приостановлено, до следующего yield:
• Последовательность Фибоначчи – это бесконечная последовательность чисел, где каждое следующее число является суммой двух предыдущих. Это классический пример задачи на генератор.

Задания для самопроверки

1. Что такое итератор и итерируемый объект?

2. Объясните, для чего нужны функции iter() и next(). Что произойдёт, если next() будет вызвана после того, как все элементы итератора будут получены?

3. Исправьте ошибку StopIteration в данном коде и объясните её:

def simple_gen():
    yield "Ты!"
    yield "Молодец!"


gen = simple_gen()
print(next(gen))
print(next(gen))
print(next(gen))

def simple_gen():
    yield "Ты!"
    yield "Молодец!"


gen = simple_gen()
print(next(gen))
print(next(gen))
#print(next(gen))

4. Напишите функцию-генератор squares_generator(n), которая принимает число n и генерирует квадраты чисел от 0 до n включительно. Используйте её для вывода на экран квадратов чисел от 0 до 5.

def squares_generator(n: int):
    for i in range(n + 1):
        yield i * i


squares = squares_generator(5)
for square in squares:
    print(square, end=" ")
# Вывод: 0 1 4 9 16 25

5. Напишите функцию-генератор get_plus(text), которая принимает строку text и при каждом вызове добавляет в конец этой строки строку "+". Трижды вызовите генератор в цикле for для строки "Отлично". Выведите на экран результат каждого вызова.

def get_plus(text: str):
    while True:
        text += "+"
        yield text


initial_text = "Отлично"
generator = get_plus(initial_text)
for i in range(3):
    print(next(generator) )
# Вывод: Отлично+
# Вывод: Отлично++
# Вывод: Отлично+++