Как устроен with

Контекстные менеджеры (оператор with) встречаются там, где перед совершением действия нужно что-то настроить, а после – прибраться. Например, чтобы прочитать файл, мы используем такой контекстный менеджер:

with open(filepath) as file_handler:
    return json.load(file_handler)

Здесь “настройкой” является открытие файла, а “уборкой” – его закрытие. Сейчас мы разберемся, как утроены контекстные менеджеры в Питоне.

Свой контекстный менеджер

Заменим нетипичную конструкцию with...as более распространенной try...finally:

file_handler = open(filepath).__enter__()
try:
    print(json.load(file_handler))
finally:
    file_handler.__exit__()

Мы видим, что файл открывается при вызове __enter__, а закрывается при вызове __exit__. При этом закрытие файла произойдет, даже если json.load бросит исключение.

Напишем собственный контекстный менеджер, который создаст книгу в Excel, а по завершении работы с ней – сохранит. Другими словами, мы хотели бы использовать его вот так:

import openpyxl


with create_workbook('workbook_name.xlsx') as workbook:
    fill_worksheet(workbook.active, some_data)

Согласно документации, для этого нам достаточно написать класс, который определяет методы __enter__ и __exit__. Вот так он будет выглядеть:

class CreateWorkbook:
    def __init__(self, filepath):
        self.filepath = filepath
    
    def __enter__(self):
        self.workbook = Workbook()
        return self.workbook

    def __exit__(self, *args):
        self.workbook.save(self.filepath)

Соответственно, использовать его можно так:

with CreateWorkbook('workbook_name.xlsx') as workbook:
    worksheet = workbook.active
    worksheet.cell(row=1, column=1, value="1")

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

Свой контекстный менеджер, ver. 2.0

Первый недостаток нашего контекстного менеджера заключается в том, что он позволяет вызвать __exit__ перед __enter__. Если это произойдет, поднимется непонятное исключение. Для того, чтобы избежать этого, воспользуемся генератором:

from openpyxl import Workbook


def create_workbook(filepath):
    workbook = Workbook()
    yield workbook
    workbook.save(filepath)


class CreateWorkbook:
    def __init__(self, filepath):
        self.filepath = filepath
    
    def __enter__(self):
        self.generator = create_workbook(self.filepath)
        return next(self.generator)

    def __exit__(self, *args):
        next(self.generator, None)

В вызове next(self.generator, None) второй аргумент нужен для того, чтобы next не поднимал исключение StopIteration.

Теперь заметим, что генератор create_workbook содержит всю предметно-ориентированную логику. Поэтому класс CreateWorkbook можно обобщить так, чтобы он от нее не зависил и работал с любыми генераторами. Заодно переименуем его в ContextManager:

class ContextManager:
    def __init__(self, generator):
        self.generator = generator
    
    def __call__(self, *args, **kwargs):
        self.args, self.kwargs = args, kwargs
        return self

    def __enter__(self):
        self.generator_instance = self.generator(*self.args, **self.kwargs)
        return next(self.generator_instance)

    def __exit__(self, *args):
        next(self.generator_instance, None)

В конструкторе ContextManager мы принимаем генератор, при вызове объекта записываем переданные генератору параметры, а при вызове __enter__ вызываем генератор. Теперь класс ContextManager можно использовать следующим образом:

from openpyxl import Workbook


def create_workbook(filepath):
    workbook = Workbook()
    yield workbook
    workbook.save(filepath)


with ContextManager(create_workbook)('workbook_name.xlsx') as workbook:
    worksheet = workbook.active
    worksheet.cell(row=1, column=1, value="1")

Вызов ContextManager(create_workbook)('workbook_name.xlsx') выглядит некрасиво. Исправим это:

create_workbook = ContextManager(create_workbook)
with create_workbook('workbook_name.xlsx') as workbook:
    worksheet = workbook.active
    worksheet.cell(row=1, column=1, value="1")

Здесь заметим, что ContextManager(create_workbook), по сути, возвращает тот же объект create_workbook, но с дополнительным поведением. Ровно для этой задачи в Питоне существует специальный синтаксический сахар – декораторы. Нам даже не потребуется изменять ContextManager:

@ContextManager
def create_workbook(filepath):
    workbook = Workbook()
    yield workbook
    workbook.save(filepath)


with create_workbook('workbook_name.xlsx') as workbook:
    worksheet = workbook.active
    worksheet.cell(row=1, column=1, value="1")

Самое приятное, что наш универсальный ContextManager уже есть в стандартной библиотеке Питона, и называется он contextmanager. Весь наш пример можно переписать так:

from contextlib import contextmanager

from openpyxl import Workbook


@contextmanager
def create_workbook(filepath):
    workbook = Workbook()
    yield workbook
    workbook.save(filepath)


with create_workbook('workbook_name.xlsx') as workbook:
    worksheet = workbook.active
    worksheet.cell(row=1, column=1, value="1")

И это еще не всё. Если после yield workbook пользовательский код поднимит исключение, мы не сохраним книгу. Чтобы этого избежать, обернем это выражение в try...finally:

@contextmanager
def create_workbook(filepath):
    workbook = Workbook()
    try:
        yield workbook
    finally:
        workbook.save(filepath)

А вот так контекстные менеджеры писать принято.

Резюме

• Контекстные менеджеры нужны там, где есть какая-то “настройка”, действия пользователя и следующая за ними “уборка”.
• У контекстного менеджера обязательно есть атрибуты __enter__ и __exit__.
• Не обязательно писать целый класс для нового контекстного менеджера, достаточно обернуть генератор в декоратор contextmanager.
• yield стоит оборачивать в блок try...finally.

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

• Более подробное описание методов __enter__ и __exit__.
• Подброрка библиотечных контекстных менеджеров и рецептов к ним.
• Часть доклада Рэймонда Хэттингера, одного из разработчиков Питона, про декораторы и контекстные менеджеры.