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

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

• __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: используйте существующую временную папку своей системы.