Python. Урок 21. Работа с контекстным менеджером

Контекстные менеджеры позволяют задать поведение при работе с конструкцией with: при входе и выходе из блока. Это упрощает работу с ресурсами в части их захвата и освобождения; транзакциями, когда нужно либо полностью закончить транзакцию, либо откатить ее целиком. Этой теме будет посвящен данный урок.

• Работа с контекстным менеджером
• Создание своего контекстного менеджера
• Работа с contextlib

Работа с контекстным менеджером

Рассмотрим пример, на котором будут показаны преимущества работы с контекстным менеджером. Задача состоит в следующем: записать в файл file.txt строку hello.

Самый простой способ ее решить – это воспользоваться функцией open() для открытия файла и записать в него данные через функцию write(). При этом нужно не забыть освободить ресурс, вызвав функцию close().

f = open(‘file.txt’, ‘w’)
f.write(‘hello’)
f.close()

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

f = open('file.txt', 'w')

try:
    f.write('hello')
except:
    print('Some error!')
finally:
    f.close()

Для того, чтобы не писать дополнительный код, связанный с обработкой исключений (это неудобно и об этом можно забыть), можно воспользоваться конструкцией with… as:

with open('file.txt', 'w') as f:
     f.write('hello')

Такая конструкция позволяет захватить ресурс (в данном случае файл), выполнить нужный набор операций (запись данных), а перед выходом – освободить ресурс.

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

Если сущность, которую вы создаете, по стилю работы с ней похожа на файл, т.е. предполагает захват ресурса и освобождение, либо требует выполнения определенных действий перед началом работы и при завершении, то хорошим решением будет создать свой контекстный менеджер, с которым можно будет работать с помощью конструкции with..as. Для этого, в класс необходимо добавить два метода: __enter__ и __exit__.

Перед тем как перейти к примеру, демонстрирующему работу с этими функциями, рассмотрим, что происходит (какие методы и в каком порядке вызываются) в конструкции:

with open('file.txt', 'w') as file_data:
    file_data.write('hello')

1. Оператор with сохраняет метод __exit__ класса File.
2. Вызывается метод __enter__ класса File.
3. __enter__ открывает файл и возвращает его.
4. Дескриптор открытого класса передается в file_data.
5. В файл записываются данные через метод write.
6. Вызывается сохраненный метод __exit__, который закрывает файл.

Если внутри конструкции with происходит исключение, то оно передается в метод __exit__, в котором производится его обработка и освобождение ресурсов (закрытие файла).

Пример реализации контекстного менеджера

Создадим класс, у объекта которого необходимо вызывать метод post_work() перед прекращением работы с ним:

class Resource:
    def __init__(self, name):
        print('Resource: create {}'.format(name))
        self.__name = name

    def get_name(self):
        return self.__name

    def post_work(self):
        print('Resource: close')

Теперь создадим контекстный менеджер для работы с Resource, который можно будет использовать с оператором with:

class ResourceForWith:
    def __init__(self, name):
        self.__resource = Resource(name)

    def __enter__(self):
        return self.__resource

    def __exit__(self, type, value, traceback):
        self.__resource.post_work()

Пример работы с ResourceForWith и конструкцией with:

with ResourceForWith('Worker') as r:
    print(r.get_name())

Если выполнить этот код, то получим следующий вывод на консоль

>python test.py
Resource: create Worker
Worker
Resource: close

Работа с contextlib

В стандартную библиотеку Python входит модуль contextlib, который содержит утилиты для построения и работы с контекстными менеджерами.

Рассмотрим только один инструмент из всего набора – contextmanager. contextmanager используется как декоратор для функции, превращая ее в контекстный менеджер. При этом схема конструирования такая: все, что написано до оператора yield вызывается в рамках функции __enter__, а все что после – в рамках __exit__.

Рассмотрим несколько примеров:

from contextlib import contextmanager

@contextmanager
def processor():
    print('--> start processing')
    yield
    print('<-- stop processing')

with processor():
    print(':: processing')

В contextmanager можно завернуть работу с файлом:

from contextlib import contextmanager

@contextmanager
def open_file(path, mode):
    f = open(path, mode)
    yield f
    f.close()

with open_file('test.txt', 'w') as fw:
    fw.write('hello')

P.S.

Вводные уроки по “Линейной алгебре на Python” вы можете найти соответствующей странице нашего сайта. Все уроки по этой теме собраны в книге “Линейная алгебра на Python”.

Если вам интересна тема анализа данных, то мы рекомендуем ознакомиться с библиотекой Pandas.  Для начала вы можете познакомиться с вводными уроками. Все уроки по библиотеке Pandas собраны в книге “Pandas. Работа с данными”.