Контекстные менеджеры

Контекстные менеджеры -- один из самых элегантных паттернов Python. Они гарантируют корректную очистку ресурсов (файлов, соединений, блокировок) даже при возникновении исключений. Инструкция with -- основной способ их использования.

Инструкция with

Базовый синтаксис with гарантирует вызов cleanup-кода:

# Without context manager — error-prone
file = open("data.txt")
try:
    content = file.read()
finally:
    file.close()

# With context manager — clean and safe
with open("data.txt") as file:
    content = file.read()
# File is automatically closed, even if an exception occurs

Несколько контекстных менеджеров

# Multiple context managers in one with statement
with (
    open("input.txt") as source,
    open("output.txt", "w") as target,
):
    for line in source:
        target.write(line.upper())

Протокол __enter__ / __exit__

Контекстный менеджер -- это объект, реализующий два метода:

class ManagedResource:
    """A resource that requires cleanup."""

    def __init__(self, name: str) -> None:
        self.name = name
        print(f"Creating resource: {name}")

    def __enter__(self):
        """Called when entering the with block."""
        print(f"Acquiring resource: {self.name}")
        return self  # This value is bound to the 'as' variable

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Called when exiting the with block (always)."""
        print(f"Releasing resource: {self.name}")
        # Return True to suppress the exception
        # Return False (or None) to propagate it
        return False

with ManagedResource("database") as resource:
    print(f"Using: {resource.name}")
    # Output:
    # Creating resource: database
    # Acquiring resource: database
    # Using: database
    # Releasing resource: database

Обработка исключений в __exit__

class SafeTransaction:
    """Database transaction context manager."""

    def __init__(self, connection) -> None:
        self.connection = connection

    def __enter__(self):
        """Start a transaction."""
        self.connection.begin()
        return self.connection

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Commit on success, rollback on error."""
        if exc_type is None:
            # No exception — commit
            self.connection.commit()
            print("Transaction committed")
        else:
            # Exception occurred — rollback
            self.connection.rollback()
            print(f"Transaction rolled back due to: {exc_val}")
        # Return False to propagate the exception
        return False

Подавление исключений

class SuppressErrors:
    """Context manager that suppresses specified exceptions."""

    def __init__(self, *exceptions: type[BaseException]) -> None:
        self.exceptions = exceptions

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Return True to suppress the exception
        if exc_type is not None and issubclass(exc_type, self.exceptions):
            print(f"Suppressed: {exc_type.__name__}: {exc_val}")
            return True
        return False

with SuppressErrors(FileNotFoundError, PermissionError):
    content = open("nonexistent.txt").read()
    # FileNotFoundError is suppressed
print("Execution continues normally")

Практические контекстные менеджеры

Timer -- измерение времени выполнения

import time

class Timer:
    """Measure execution time of a code block."""

    def __init__(self, label: str = "Block") -> None:
        self.label = label
        self.start: float = 0
        self.elapsed: float = 0

    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start
        print(f"{self.label} took {self.elapsed:.4f}s")
        return False

with Timer("Data processing") as t:
    data = [i ** 2 for i in range(1_000_000)]

print(f"Elapsed: {t.elapsed:.4f}s")

TemporaryDirectory -- временная директория

import os
import shutil

class TemporaryWorkDir:
    """Create and cleanup a temporary working directory."""

    def __init__(self, base_path: str = "/tmp") -> None:
        self.path = os.path.join(base_path, f"work_{os.getpid()}")

    def __enter__(self) -> str:
        os.makedirs(self.path, exist_ok=True)
        return self.path

    def __exit__(self, exc_type, exc_val, exc_tb):
        shutil.rmtree(self.path, ignore_errors=True)
        return False

with TemporaryWorkDir() as workdir:
    # Work with temporary files
    filepath = os.path.join(workdir, "temp.txt")
    with open(filepath, "w") as f:
        f.write("temporary data")
# Directory is cleaned up automatically

Модуль contextlib

Стандартная библиотека предоставляет мощные инструменты для создания контекстных менеджеров.

@contextmanager -- из генератора

Декоратор @contextmanager превращает генератор в контекстный менеджер:

from contextlib import contextmanager

@contextmanager
def managed_connection(host: str, port: int):
    """Create a database connection with automatic cleanup."""
    print(f"Connecting to {host}:{port}")
    connection = {"host": host, "port": port, "active": True}

    try:
        yield connection  # Value bound to 'as' variable
    except Exception as e:
        print(f"Error during connection use: {e}")
        raise
    finally:
        # Cleanup — always runs
        connection["active"] = False
        print(f"Disconnected from {host}:{port}")

with managed_connection("localhost", 5432) as conn:
    print(f"Using connection: {conn}")

@contextmanager -- реальные примеры

from contextlib import contextmanager
import os
import sys
from io import StringIO

@contextmanager
def change_directory(path: str):
    """Temporarily change the working directory."""
    original = os.getcwd()
    try:
        os.chdir(path)
        yield path
    finally:
        os.chdir(original)

@contextmanager
def capture_output():
    """Capture stdout and stderr."""
    old_stdout, old_stderr = sys.stdout, sys.stderr
    sys.stdout = StringIO()
    sys.stderr = StringIO()
    try:
        yield sys.stdout, sys.stderr
    finally:
        sys.stdout, sys.stderr = old_stdout, old_stderr

@contextmanager
def environment_variable(key: str, value: str):
    """Temporarily set an environment variable."""
    old_value = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if old_value is None:
            del os.environ[key]
        else:
            os.environ[key] = old_value

# Usage
with environment_variable("DEBUG", "true"):
    print(os.environ["DEBUG"])  # "true"
# Original value is restored

contextlib.suppress

Элегантная замена try/except/pass:

from contextlib import suppress

# Instead of:
try:
    os.remove("temp.txt")
except FileNotFoundError:
    pass

# Use:
with suppress(FileNotFoundError):
    os.remove("temp.txt")

# Multiple exception types
with suppress(FileNotFoundError, PermissionError):
    os.remove("protected.txt")

contextlib.ExitStack

ExitStack управляет динамическим набором контекстных менеджеров:

from contextlib import ExitStack

def process_multiple_files(paths: list[str]) -> list[str]:
    """Open and process multiple files safely."""
    with ExitStack() as stack:
        # Open all files — all will be closed on exit
        files = [stack.enter_context(open(path)) for path in paths]

        # Process all files
        results = []
        for f in files:
            results.append(f.read())
        return results

# Dynamic number of context managers
def create_connections(hosts: list[str]) -> list[dict]:
    """Create multiple connections with guaranteed cleanup."""
    with ExitStack() as stack:
        connections = []
        for host in hosts:
            conn = stack.enter_context(managed_connection(host, 5432))
            connections.append(conn)

        # Use all connections
        return [{"host": c["host"], "active": c["active"]} for c in connections]

ExitStack с callback

from contextlib import ExitStack

def setup_environment():
    """Set up environment with cleanup callbacks."""
    with ExitStack() as stack:
        # Register cleanup callbacks
        stack.callback(print, "Cleanup step 3: done")
        stack.callback(print, "Cleanup step 2: done")
        stack.callback(print, "Cleanup step 1: done")

        print("Environment is set up")
        # On exit, callbacks run in LIFO order:
        # Cleanup step 1: done
        # Cleanup step 2: done
        # Cleanup step 3: done

Async-контекстные менеджеры

Для асинхронного кода используются async with и методы __aenter__/__aexit__:

import asyncio
from contextlib import asynccontextmanager

class AsyncDatabasePool:
    """Async database connection pool."""

    async def __aenter__(self):
        """Asynchronously acquire a connection."""
        print("Acquiring async connection...")
        await asyncio.sleep(0.1)  # Simulate async connection
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Release the connection."""
        print("Releasing async connection...")
        await asyncio.sleep(0.05)  # Simulate async cleanup
        return False

    async def query(self, sql: str) -> list:
        """Execute a query."""
        await asyncio.sleep(0.01)
        return [{"result": sql}]

async def main():
    async with AsyncDatabasePool() as pool:
        result = await pool.query("SELECT * FROM users")
        print(result)

@asynccontextmanager

from contextlib import asynccontextmanager
import asyncio

@asynccontextmanager
async def async_timer(label: str):
    """Async version of the timer context manager."""
    import time
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed:.4f}s")

async def main():
    async with async_timer("API call"):
        await asyncio.sleep(0.5)
    # Output: API call: 0.5004s

Паттерн: контекстный менеджер как декоратор

contextlib.ContextDecorator позволяет использовать контекстный менеджер и как with, и как декоратор:

from contextlib import ContextDecorator
import time

class log_execution(ContextDecorator):
    """Log function execution time and exceptions."""

    def __init__(self, label: str = "block") -> None:
        self.label = label

    def __enter__(self):
        self.start = time.perf_counter()
        print(f"[{self.label}] Starting...")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        elapsed = time.perf_counter() - self.start
        if exc_type:
            print(f"[{self.label}] Failed after {elapsed:.4f}s: {exc_val}")
        else:
            print(f"[{self.label}] Completed in {elapsed:.4f}s")
        return False

# Use as context manager
with log_execution("data processing"):
    data = [i ** 2 for i in range(100_000)]

# Use as decorator
@log_execution("compute")
def heavy_computation():
    return sum(i ** 2 for i in range(1_000_000))

heavy_computation()