Можно ли создать асинхронный контекстный менеджер в Python?

# Асинхронные контекстные менеджеры в Python

Короткий ответ: ДА, можно и нужно!

Полностью поддерживаются через async with и @asynccontextmanager декоратор (или класс с методами __aenter__ и __aexit__).

1. Класс с методами __aenter__ и __aexit__

Это асинхронный аналог обычного контекстного менеджера:

import asyncio

class AsyncDatabaseConnection:
    """Асинхронный контекстный менеджер для работы с БД"""
    
    def __init__(self, connection_string: str):
        self.connection_string = connection_string
        self.connection = None
    
    async def __aenter__(self):
        """Вызывается при входе в блок async with"""
        print(f"Подключение к {self.connection_string}...")
        # Имитируем асинхронное подключение
        await asyncio.sleep(0.5)
        self.connection = f"Connection to {self.connection_string}"
        print("Подключено!")
        return self.connection  # Это будет значение после as
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Вызывается при выходе из блока async with"""
        if exc_type:
            print(f"Ошибка: {exc_type.__name__}: {exc_val}")
        print("Отключение...")
        await asyncio.sleep(0.3)  # Имитируем закрытие соединения
        self.connection = None
        print("Отключено!")
        return False  # False = не подавляем исключение

# Использование:
async def main():
    async with AsyncDatabaseConnection("postgresql://localhost/mydb") as conn:
        print(f"Используем: {conn}")
        # Здесь можно выполнять операции с БД
        await asyncio.sleep(1)
    # После блока __aexit__ уже вызван

asyncio.run(main())

2. Декоратор @asynccontextmanager

Более элегантный способ (из модуля contextlib):

from contextlib import asynccontextmanager
import asyncio

@asynccontextmanager
async def async_db_connection(db_url: str):
    """Асинхронный контекстный менеджер как функция"""
    print(f"Подключение к {db_url}...")
    await asyncio.sleep(0.5)  # Асинхронное подключение
    conn = f"Connection: {db_url}"
    
    try:
        print("Подключено!")
        yield conn  # Передаём соединение пользователю
    except Exception as e:
        print(f"Ошибка при работе: {e}")
        raise
    finally:
        print("Закрываем соединение...")
        await asyncio.sleep(0.3)
        print("Соединение закрыто!")

# Использование:
async def main():
    async with async_db_connection("postgresql://localhost/mydb") as conn:
        print(f"Используем: {conn}")
        # Работаем с БД
        await asyncio.sleep(1)

asyncio.run(main())

3. Практический пример: работа с Redis

from contextlib import asynccontextmanager
import aioredis
import asyncio

class RedisManager:
    """Менеджер подключений к Redis с пулом"""
    
    def __init__(self, redis_url: str):
        self.redis_url = redis_url
        self.redis = None
    
    async def __aenter__(self):
        self.redis = await aioredis.create_redis_pool(self.redis_url)
        return self.redis
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.redis:
            self.redis.close()
            await self.redis.wait_closed()
        if exc_type:
            print(f"Ошибка: {exc_val}")
        return False

# Или через декоратор:
@asynccontextmanager
async def get_redis_connection(redis_url: str):
    redis = await aioredis.create_redis_pool(redis_url)
    try:
        yield redis
    finally:
        redis.close()
        await redis.wait_closed()

# Использование в асинхронной функции:
async def cache_data():
    async with get_redis_connection("redis://localhost") as redis:
        await redis.set("key", "value")
        value = await redis.get("key")
        print(f"Из кэша: {value}")

4. Вложенные асинхронные контекстные менеджеры

@asynccontextmanager
async def timer():
    """Контекстный менеджер для измерения времени"""
    import time
    start = time.time()
    try:
        yield
    finally:
        elapsed = time.time() - start
        print(f"Выполнено за {elapsed:.2f} секунд")

@asynccontextmanager
async def database_transaction():
    """Контекстный менеджер для транзакций"""
    print("BEGIN TRANSACTION")
    try:
        yield
        print("COMMIT")
    except Exception:
        print("ROLLBACK")
        raise

async def main():
    async with timer():
        async with database_transaction():
            await asyncio.sleep(1)
            print("Выполняем операции...")
            # Можно вызвать исключение для проверки ROLLBACK

asyncio.run(main())

5. Важные различия от обычных контекстных менеджеров

6. Обработка исключений

@asynccontextmanager
async def handle_errors():
    try:
        yield
    except ValueError as e:
        print(f"ValueError: {e}")
        # Можешь либо подавить (не raise), либо пробросить дальше
        raise  # Пробрасываем исключение выше
    except Exception as e:
        print(f"Неожиданная ошибка: {e}")
        raise
    finally:
        print("Cleanup код выполняется в любом случае")

async def main():
    async with handle_errors():
        raise ValueError("Что-то пошло не так!")

asyncio.run(main())

Реальный примеры из фреймворков

FastAPI (работа с БД сессией)

from fastapi import FastAPI
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Код инициализации (startup)
    engine = create_async_engine("postgresql+asyncpg://...")
    print("Приложение запустилось")
    yield
    # Код очистки (shutdown)
    await engine.dispose()
    print("Приложение завершилось")

app = FastAPI(lifespan=lifespan)

Ключевые преимущества

1. Безопасность ресурсов — гарантирует закрытие соединений даже при ошибках
2. Читаемость — код понятнее, чем try-finally
3. Переиспользуемость — логику инициализации/очистки можно переиспользовать
4. Асинхронность — не блокируешь событийный цикл во время ожидания

Вывод

Асинхронные контекстные менеджеры — это мощный инструмент для управления ресурсами в асинхронном коде. Используй их везде, где нужна инициализация и очистка ресурсов (БД, сеть, файлы, пулы соединений).