python
August 17, 2022

Как автоматизировать операции с контекстом Python

Часто участки кода требуется дополнять одними и теми же инструкциями. Например, для добавления времени исполнения строк или автоматического открытия и закрытия файла. Эти действия можно автоматизировать при помощи декоратора contextmanager из модуля contextlib.
Техника его применения заключается в декорировании некоторой функции, в которой выполняются необходимые вам повторяющиеся действия, а точкой входа для изменяемого кода является строка с инструкцией yield. Например, так можно создать контекст измерения времени:

import numpy as np
from time import time
from contextlib import contextmanager

@contextmanager
def code_time():
    start = time()
    yield
    print(time() - start)

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

with code_time():
    ar1 = np.random.normal(0, 3, size=(1000,1000))
    ar2 = np.random.normal(-2, 3, size=(1000,1000))
    sum = ar1 + ar2

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

@contextmanager
def writefile(fn):   
    try:
        f = open(fn,'w')
        yield f
    finally:
        f.close()

А теперь запишем строку в файл новым способом:

with writefile('temp.txt') as f:
    f.write('Привет')