Практические шаблоны менеджеров контекста

Введение

Конвейер обработки файлов

import io
src = io.StringIO('line1\nline2\nline3')
dst = io.StringIO()
with src, dst:
    for line in src:
        dst.write(line.upper())
print(dst.getvalue())

Транзакция базы данных

from contextlib import contextmanager
@contextmanager
def db_transaction(conn):
    try:
        yield conn
        print('COMMIT')
    except Exception as e:
        print(f'ROLLBACK: {e}')
        raise
    finally:
        print('CLOSE')
with db_transaction({'execute': lambda q: print(f'SQL: {q}')}) as db:
    db['execute']('INSERT ...')

Блокировка потоков

import threading
lock = threading.Lock()
shared = []
def add(x):
    with lock:
        shared.append(x)
add(1); add(2)
print(shared)

Смена каталога

from contextlib import contextmanager
import os
@contextmanager
def chdir(path):
    old = os.getcwd()
    os.chdir(path)
    try: yield
    finally: os.chdir(old)
original = os.getcwd()
with chdir('/tmp'):
    print('in tmp:', os.getcwd())
print('back:', os.getcwd())

Временные файлы

import tempfile
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=True) as f:
    f.write('temp data')
    print('temp file:', f.name)
print('deleted')

Изменение переменной окружения

from contextlib import contextmanager
import os
@contextmanager
def env(**kwargs):
    old = {k: os.environ.get(k) for k in kwargs}
    os.environ.update({k:str(v) for k,v in kwargs.items()})
    try: yield
    finally:
        for k,v in old.items():
            if v is None: os.environ.pop(k,None)
            else: os.environ[k]=v
with env(DEBUG='1', LOG_LEVEL='DEBUG'):
    print(os.environ['DEBUG'])

Блок профилирования

from contextlib import contextmanager
import cProfile
@contextmanager
def profile():
    pr = cProfile.Profile()
    pr.enable()
    yield pr
    pr.disable()
with profile() as pr:
    sum(range(1000000))
print('profiled')

Сетевой сокет

import socket
try:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(0.1)
        print('socket ready')
except Exception as e:
    print(f'socket: {e}')

Перехват предупреждений

import warnings
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter('always')
    warnings.warn('test warning')
print(len(w), w[0].category.__name__)

Комбинирование шаблонов

from contextlib import contextmanager
import time
@contextmanager
def timed_operation(name):
    print(f'[{name}] start')
    t0 = time.perf_counter()
    try:
        yield
        print(f'[{name}] done in {time.perf_counter()-t0:.4f}s')
    except Exception as e:
        print(f'[{name}] failed: {e}')
        raise
with timed_operation('loop'):
    sum(range(100000))

Быстрая проверка

Итоги

Продолжайте