Создание своего менеджера контекста

Чтобы создать собственный контекстный менеджер, класс должен определить два метода:

• __enter__(self) – код, выполняющийся при входе в блок with. Его результат будет привязан к имени после as или проигнорирован, если as не используется.

• __exit__(self, exc_type, exc_value, exc_traceback) – код, выполняющийся при выходе из блока with. Здесь можно осуществить завершение/очистку ресурса. Если нужно подавить возникшее исключение, этот метод должен вернуть True. Обычно же возвращается False (или ничего), чтобы исключение далее было возбуждено.

Рассмотрим пример: создадим менеджер контекста, который измеряет время выполнения блока кода.

import time

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self  # вернем сам объект, чтобы им можно было пользоваться внутри блока

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.end = time.perf_counter()
        self.interval = self.end - self.start
        print(f"Блок кода выполнялся {self.interval:.4f} секунд")
        # Не подавляем исключения, поэтому возвращаем False (по умолчанию None тоже трактуется как False).

# Использование:
with Timer() as t:
    # какой-то код
    total = 0
    for i in range(1000000):
        total += i*i
print(f"Время, замеренное таймером: {t.interval:.4f} секунд")

Условный пример вывода; фактическое время зависит от окружения:

Блок кода выполнялся 0.1234 секунд
Время, замеренное таймером: 0.1234 секунд

Здесь класс Timer в методе __enter__ сохраняет текущее время self.start и возвращает себя. В методе __exit__ вычисляет интервал времени и выводит его на экран. Также сохраняет интервал в self.interval, чтобы мы могли после блока with обратиться к этому значению (например, для дополнительной обработки). Мы решили не подавлять возможное исключение внутри блока (метод возвращает None), чтобы оно при необходимости было выброшено наружу.

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

import os

class ChangeDir:
    def __init__(self, new_path):
        self.new_path = new_path
        self.old_path = None

    def __enter__(self):
        self.old_path = os.getcwd()  # запоминаем текущую директорию
        os.chdir(self.new_path)      # переходим в новую
        return None  # можем не возвращать ничего (или вернуть self, если нужно)

    def __exit__(self, exc_type, exc_val, exc_tb):
        os.chdir(self.old_path)      # возвращаемся в старую директорию
        # Исключения не обрабатываем специально, возвращаем False

Использование:

print("Before:", os.getcwd())
with ChangeDir("/tmp"):
    # внутри блока текущая директория изменена на /tmp
    print("Inside:", os.getcwd())
# после выхода директория восстановлена
print("After:", os.getcwd())

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

os.chdir меняет рабочую директорию всего процесса. Этот пример предназначен для последовательного кода; он не подходит для параллельных потоков или корутин, зависящих от рабочей директории. Путь /tmp — пример для Unix: используйте существующую временную папку своей системы.

FAQ

Задавайте вопросы в комментариях — мы будем расширять раздел.

1. Зачем __enter__ возвращает self (или другой объект)?
   Возвращённое значение привязывается к имени после as ... в конструкции with. Если нужен доступ к данным менеджера внутри блока, обычно возвращают self (with Timer() as t:). Если ничего не нужно — можно вернуть None.
2. Что будет, если в __exit__ вернуть True?
   Исключение, возникшее в блоке with, будет подавлено и не «уйдёт» выше по стеку. Возврат False (или None) оставляет исключение активным.
3. Можно ли использовать print в __exit__ вместо логирования?
   Технически можно, но для реального кода лучше применить logging — так сообщения можно фильтровать и перенаправлять.
4. Что произойдёт, если в __exit__ само возникнет новое исключение?
   Наружу выйдет новое исключение, а исходное обычно останется в его __context__ и будет видно в цепочке traceback. contextlib.ExitStack сам по себе не подавляет ошибки очистки. Политику обработки таких ошибок нужно задать явно.
5. Нужно ли явно вызывать Timer.__exit__ вне with?
   После успешного __enter__ конструкция with вызовет __exit__ автоматически. Ручной вызов допустим (например, при тестировании), но редко нужен.
6. Как сделать менеджер, который только подавляет исключения определённого типа?
   Проверяйте exc_type и возвращайте True только для нужных исключений:

   
   def __exit__(self, exc_type, exc_val, tb):
       return exc_type is FileNotFoundError
       

7. Можно ли написать менеджер без класса?
   Да, с помощью @contextlib.contextmanager:

   
   import time
   from contextlib import contextmanager
   @contextmanager
   def timer():
           start = time.perf_counter()
           try:
               yield
           finally:
               print(time.perf_counter() - start)
       

   Код до yield выполняется при входе. Исключение из блока возобновляет генератор в точке yield; очистка в finally выполняется и при таком выходе.
8. Будет ли ChangeDir безопасен при вложенных вызовах?
   Если одновременно используется несколько экземпляров («матрёшка»), каждый запоминает свою old_path, поэтому восстановление происходит корректно в последовательном коде, если исходные директории доступны. Повторный вложенный вход в тот же экземпляр перезапишет old_path. Изменение рабочей директории действует на весь процесс, поэтому параллельное использование небезопасно.