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

Управление ресурсами

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

Для работы с файлами часто используется конструкция, показанная в примере ниже. Она позволяет захватить ресурс, выполнить необходимые операции, а затем освободить его:

with open("test.txt") as f:    
    data = f.read()

Давайте возьмем пример управления файлами. Когда файл открывается, используется дескриптор файла, который является ограниченным ресурсом. Только определенное количество файлов может быть открыто процессом одновременно. Следующая программа демонстрирует это.

file_descriptors = [] 
for x in range(100000): 
    file_descriptors.append(open('test.txt', 'w'))

Результат:

Traceback (most recent call last):
  File "context.py", line 3, in
OSError: [Errno 24] Too many open files: 'test.txt'

Мы получили сообщение об ошибке, указывающее, что открыто слишком много файлов. Приведенный выше пример представляет собой случай утечки файлового дескриптора. Это происходит потому, что слишком много открытых файлов и они не закрыты. Бывает, что программист просто забывает закрыть открытый файл.

Управление ресурсами с помощью контекстного менеджера

Предположим, что блок кода вызывает исключение или содержит сложный алгоритм с несколькими инструкциями return, и тогда становится неудобно закрывать файл во всех местах.

Как правило, в других языках при работе с файлами для того, чтобы файловый ресурс был закрыт после использования, даже если было возбуждено исключение, используется  конструкция try-except-finally. Python предоставляет простой способ управления ресурсами: контекстные менеджеры. Выражение, следующее после ключевого слова with, должно возвращать объект, реализующий протокол Context Manager. Контекстные менеджеры могут быть написаны с использованием классов или функций (с помощью декораторов).

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

При создании контекстных менеджеров с использованием классов пользователь должен убедиться, что у класса есть следующие методы:  __enter__() и __exit__(). __enter__() возвращает требуемый ресурс, а __exit __() ничего не возвращает, но выполняет освобождение ресурсов.

Для начала, давайте создадим простой класс ContextManager, чтобы понять базовую структуру создания контекстных менеджеров с использованием классов, как показано ниже:

class ContextManager(): 
    def __init__(self): 
        print('init method called') 
          
    def __enter__(self): 
        print('enter method called') 
        return self
      
    def __exit__(self, exc_type, exc_value, exc_traceback): 
        print('exit method called') 
  
  
with ContextManager() as manager: 
    print('with statement block')

Результат:

init method called
enter method called
with statement block
exit method called

В этом случае создается объект ContextManager. Он присваивается переменной, следующей за ключевым словом as (manager в данном случае).  При запуске вышеуказанной программы последовательно выполняются следующие действия:

  • __init__()
  • __enter__()
  • инструкции (код внутри блока with)
  • __exit__() [параметры этого метода используются для управления исключениями]

Управление файлами с помощью контекстного менеджера

Давайте используем вышеприведенную концепцию для создания класса, который применяется для управлении файловыми ресурсами. Класс FileManager помогает открыть файл, записать/прочитать содержимое, а затем закрыть его.

class FileManager(): 
    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_value, exc_traceback): 
        self.file.close() 
  
# загрузка файла
with FileManager('test.txt', 'w') as f: 
    f.write('Test') 
  
print(f.closed)

Результат:

True

Управление файлами с помощью контекстного менеджера и конструкции with

При выполнении блока with последовательно выполняются следующие операции:

  • Создается объект FileManager с аргументами test.txt и w (write) при выполнении метода __init__.
  • Метод __enter__ открывает файл test.txt в режиме записи (write) и возвращает объект FileManager, который присваивается переменной f.
  • Текст «Test» записывается в файл.
  • Метод __exit__ обеспечивает закрытие файла при выходе из блока with.

Когда выполняется print(f.closed), то на экран выводится True, поскольку FileManager уже позаботился о закрытии файла, что в противном случае нужно было бы сделать явно.

Управление подключением к базе данных с помощью контекстного менеджера

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

from pymongo import MongoClient 
  
class MongoDBConnectionManager(): 
    def __init__(self, hostname, port): 
        self.hostname = hostname 
        self.port = port 
        self.connection = None
  
    def __enter__(self): 
        self.connection = MongoClient(self.hostname, self.port) 
        return self
  
    def __exit__(self, exc_type, exc_value, exc_traceback): 
        self.connection.close() 
  
# подключение к localhost
with MongoDBConnectionManager('localhost', '27017') as mongo: 
    collection = mongo.connection.SampleDb.test 
    data = collection.find({'_id': 1}) 
    print(data.get('name'))

Управление подключениями к базе данных с помощью контекстного менеджера и оператора with

При выполнении блока with последовательно выполняются следующие операции:

  • Создается объект MongoDBConnectionManager с аргументами localhost и 27017 при выполнении метода __ init__.
  • Метод __enter__ открывает соединение mongodb и возвращает объект MongoDBConnectionManager переменной mongo.
  • Осуществляется доступ к коллекции (collection) test в базе данных SampleDb и извлекается документ с _id=1. На экран выводится поле name документа.
  • Метод __exit__ обеспечивает закрытие соединения при выходе из блока with.