Как можно кастомизировать поведение контекстного менеджера для разных типов ресурсов (например, временные файлы, базы данных, соединения)?

Существует несколько способов кастомизации поведения контекстного менеджера для разных типов ресурсов в Python, обеспечивая гибкость и повторное использование кода. Основная идея заключается в создании классов, которые реализуют протокол контекстного менеджера (методы __enter__ и __exit__), и параметризовать эти классы для работы с разными типами ресурсов.

1. Параметризация класса контекстного менеджера: Наиболее распространенный подход - это передача информации о типе ресурса и необходимых параметрах при создании экземпляра класса контекстного менеджера.

class GenericResourceContextManager:
    def __init__(self, resource_type, *args, **kwargs):
        self.resource_type = resource_type
        self.args = args
        self.kwargs = kwargs
        self.resource = None

    def __enter__(self):
        if self.resource_type == 'file':
            self.resource = open(*self.args, **self.kwargs) # Пример открытия файла
        elif self.resource_type == 'database':
            # Логика подключения к базе данных, используя self.args и self.kwargs
            # Например, self.resource = psycopg2.connect(*self.args, **self.kwargs)
            print("Подключение к базе данных (имитация)")
            self.resource = "Database Connection" #Имитация соединения
        elif self.resource_type == 'connection':
           # Логика создания соединения, например, сетевого
           print("Создание соединения (имитация)")
           self.resource = "Network Connection" # Имитация соединения
        else:
            raise ValueError(f"Неизвестный тип ресурса: {self.resource_type}")
        return self.resource

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.resource_type == 'file':
            self.resource.close()
        elif self.resource_type == 'database':
            # Логика закрытия соединения с базой данных
            print("Закрытие соединения с базой данных (имитация)")
        elif self.resource_type == 'connection':
            # Логика закрытия соединения
            print("Закрытие соединения (имитация)")
        self.resource = None
        return False  #  False чтобы не подавлять исключения, если они были

# Пример использования
with GenericResourceContextManager('file', 'temp.txt', 'w') as f:
    f.write("Привет, мир!")

with GenericResourceContextManager('database', database='mydatabase', user='myuser') as db:
    # Выполнение операций с базой данных
    print(f"Работа с ресурсом: {db}")

with GenericResourceContextManager('connection') as conn:
    print(f"Работа с соединением: {conn}")


    

2. Использование наследования: Можно создать базовый класс контекстного менеджера и затем наследовать от него для каждого типа ресурса, переопределяя методы __enter__ и __exit__.

class BaseContextManager:
    def __enter__(self):
        raise NotImplementedError

    def __exit__(self, exc_type, exc_val, exc_tb):
        raise NotImplementedError

class FileContextManager(BaseContextManager):
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
        self.file = None

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.file:
            self.file.close()
        return False #  False чтобы не подавлять исключения, если они были

# Пример использования
with FileContextManager('my_file.txt', 'w') as f:
    f.write("Текст в файле")
    

3. Использование фабричных функций: Можно создать фабричную функцию, которая возвращает экземпляр контекстного менеджера, настроенный для конкретного типа ресурса.

def context_manager_factory(resource_type, *args, **kwargs):
    if resource_type == 'file':
        class FileContextManager:
            def __init__(self, filename, mode):
                self.filename = filename
                self.mode = mode
                self.file = None

            def __enter__(self):
                self.file = open(self.filename, self.mode)
                return self.file

            def __exit__(self, exc_type, exc_val, exc_tb):
                if self.file:
                    self.file.close()
                return False

        return FileContextManager(*args, **kwargs)
    elif resource_type == 'database':
        # Здесь логика создания контекстного менеджера для базы данных
        print("Создание контекстного менеджера для базы данных (фабрика)")
        class DBContextManager:
          def __init__(self, db_name):
            self.db_name = db_name

          def __enter__(self):
            print(f"Подключение к БД: {self.db_name}")
            return "DB Connection"

          def __exit__(self, exc_type, exc_val, exc_tb):
            print(f"Закрытие соединения с БД: {self.db_name}")
            return False

        return DBContextManager(*args, **kwargs)

    else:
        raise ValueError("Неподдерживаемый тип ресурса")

# Пример использования
with context_manager_factory('file', 'data.txt', 'w') as file:
    file.write("Более сложный пример с фабрикой")

with context_manager_factory('database', 'mydb') as db:
  print(f"Работа с базой данных: {db}")

    

4. Использование сторонних библиотек: Некоторые библиотеки (например, contextlib) предоставляют инструменты для создания и кастомизации контекстных менеджеров, упрощая разработку. Например, можно использовать contextlib.contextmanager в сочетании с генераторами для создания простых контекстных менеджеров.

import contextlib

@contextlib.contextmanager
def database_connection(database_url):
    connection = None
    try:
        # Логика подключения к базе данных
        print(f"Подключение к базе данных {database_url}...")
        connection = "Database Connection"  #Имитация подключения
        yield connection
    finally:
        if connection:
            # Логика закрытия соединения с базой данных
            print("Закрытие соединения с базой данных...")

# Пример использования
with database_connection("postgresql://user:password@host:port/database") as db:
    print(f"Работа с {db}")
    

Выбор конкретного подхода зависит от сложности задачи и требуемой степени гибкости. Параметризация класса – хороший выбор для простых случаев. Наследование обеспечивает большую структуру и контроль над поведением каждого типа ресурса. Фабричные функции предлагают динамическое создание контекстных менеджеров. Библиотека contextlib предоставляет удобные инструменты для создания простых контекстных менеджеров с помощью генераторов.