Как устроены генераторы в Python

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

Пара слов об итераторах

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

class ReverseIt():
    def __init__(self, reverse_me):
        self.reverse_me = reverse_me
        self.current_index = len(reverse_me) - 1

    def __iter__(self):
        return self

    def __next__(self):
        if self.current_index < 0:
            raise StopIteration()
        current_element = self.reverse_me[self.current_index]
        self.current_index -= 1
        return current_element

Метод __iter__ позволяет использовать итератор в конструкции for ... in ..., а __next__ определяет порядок обхода коллекции. Этих двух методов достаточно, чтобы класс удовлетворял протоколу итератора. Заметим, что итератор не обязан работать с коллекцией. Метод __next__ может возвращать любое значение и в принципе делать что угодно. Например, можно заставить его загружать данные с очередной страницы сайта. Или исполнять __enter__ при первом вызове и __exit__ при втором. Или выполнять только часть всех вычислений. Идея о том, что __next__ может не только задавать порядок обхода последовательности, но и выполнять произвольный код, легла в основу генераторов.

Устройство генератора

Генератор – это объект такого класса, который удовлетворяет протоколу итератора. Но генератор, в отличие от итератора, не обходит (не итерирует) коллекцию, а возвращает (генерирует) новые данные. Вот так выглядит генератор, который постранично загружает данные о попытках сдать задачи на devman.org:

import requests


class SolutionAttemptsFetcher():
    SOLUTION_ATTEMPTS_URL = 'https://devman.org/api/challenges/solution_attempts/'
    PAGE_LIMIT = 10

    def __init__(self):
        self.current_page_number = 1

    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current_page_number >= SolutionAttemptsFetcher.PAGE_LIMIT:
            raise StopIteration()
        params = {'page': self.current_page_number}
        page = requests.get(SolutionAttemptsFetcher.SOLUTION_ATTEMPTS_URL, params).json()
        self.current_page_number += 1
        return page['records']

Мы его можем использовать точно так же, как и итератор:

>>> attempt_pages = SolutionAttemptsFetcher()
>>> for attempts in attempt_pages:
...  print(len(attempts))
...  
30
30
30
30
30
30
30
30
30

Перед выводом длины очередной страницы происходит заметная задержка. Она вызвана тем, что метод __next__ взаимодействует с сетью и парсит JSON, а это долгие операции.

Генератор как функция

Питон дает записать генератор в виде функции. Для этого используется ключевое слово yield. Так выглядит SolutionAttemptsFetcher, если его написать как функцию:

def fetch_solution_attempts():
    solution_attempts_url = 'https://devman.org/api/challenges/solution_attempts/'
    page_limit = 10
    for current_page_number in range(1, page_limit):
        params = {'page': current_page_number}
        page = requests.get(solution_attempts_url, params).json()
        yield page['records']

Проще всего думать о yield как о таком return, который ставит функцию на паузу, чтобы при следующем вызове она продолжила работу именно с этого места. Другими словами, следующий фрагмент кода выведет сначала 1, а потом 2:

def count_to_two():
    yield 1
    yield 2

counter = count_to_two()
print(next(counter))
print(next(counter))

Generator expressions

Generator expression работает и выглядит очень похоже на list comprehension. Перепишем наш пример в этом виде:

def fetch_solution_attempts_page(page_number):
    solution_attempts_url = 'https://devman.org/api/challenges/solution_attempts/'
    params = {'page': page_number}
    page = requests.get(solution_attempts_url, params).json()
    return page['records']

page_limit = 10
attempt_pages = (fetch_solution_attempts_page(current_page_number) for current_page_number in range(1, page_limit))
for attempt_page in attempt_pages:
    print(len(attempt_page))

Вместо круглых скобок мы могли бы написать квадратные и получить список страниц в переменной attempt_pages. Но так как мы используем круглые скобки, в attempt_pages кладется генератор, который вызывает функцию fetch_solution_attempts_page и возвращает результат ее работы при вызове next(attempt_pages). Выходит, мы не сохраняем за раз больше одной страницы. Это эффективнее, чем list comprehension, если нам надо обойти все страницы один раз и этих страниц потенциально может быть много.

Ничем, кроме записи, generator expressions не отличаются от показанных выше генераторов.

Зачем используют генераторы

1. Ленивые вычисления.

Предположим, нам надо найти последнюю отправку человека с каким-то юзернеймом. Это легко сделать с помощью генератора написанного нами fetch_solution_attempts, который вместо загрузки десятка страниц будет загружать их по одной. Так, если мы найдем нужную запись на странице 4, мы сэкономим время на загрузке и обработке еще шести страниц.

2. Переплетение клиентского и библиотечного кода.

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

3. Соблюдение порядка вызовов.

Когда у нас есть функция func2, которая должна быть вызвана после клиенсткого кода, который зависит от func1, мы можем закрепить порядок их вызовов в генераторе:

def api():
    yield func1()
    yield func2()

Резюме

• Генераторы – это почти итераторы.
• Чтобы написать генератор, достаточно использовать ключевое слово yield вместо return.
• Генераторы часто используются, чтобы обрабатывать результаты выполнения долгой функции порционно.

Дальнейшее чтение

• Документация к itertools, которые построены на генераторах.
• Часть DiveIntoPython3 про то, как решать числовые ребусы с помощью itertools.
• Статья про ленивые вычисления на Википедии.
• PEP 255, предложивший введение генераторов.