Python contextlib模块高级使用技巧

Python contextlib模块高级使用技巧

contextlib提供了比with语句更丰富的上下文管理工具。@contextmanager装饰器将一个生成器函数转为上下文管理器：

from contextlib import contextmanager

@contextmanager
def managed_resource(*args, **kwargs):
resource = acquire_resource(*args, **kwargs)
try:
yield resource
finally:
release_resource(resource)

yield之前的代码对应__enter__，yield之后的对应__exit__。finally确保异常时也能清理。

contextmanager的实现原理：

class _GeneratorContextManager:
def __init__(self, func, args, kwds):
self.gen = func(*args, **kwds)
self.func, self.args, self.kwds = func, args, kwds

def __enter__(self):
try:
return next(self.gen)
except StopIteration:
raise RuntimeError("generator didn't yield")

def __exit__(self, type, value, traceback):
if type is None:
try:
next(self.gen)
except StopIteration:
return False
else:
raise RuntimeError("generator didn't stop")
else:
if value is None:
value = type()
try:
self.gen.throw(type, value, traceback)
except StopIteration as exc:
return exc is not value
except RuntimeError as exc:
if exc is value:
return False
if exc.__cause__ is value:
return False
raise
except BaseException:
if sys.exc_info()[1] is value:
return False
raise
return True

__enter__调用next启动生成器到yield。__exit__恢复生成器执行清理代码。异常通过throw传递到生成器内部。

ExitStack组合管理多个上下文管理器：

from contextlib import ExitStack

def clean_up_resources():
stack = ExitStack()
try:
file1 = stack.enter_context(open('file1.txt'))
file2 = stack.enter_context(open('file2.txt'))
db = stack.enter_context(db_connection())
return process_files(file1, file2, db)
except:
stack.close()
raise

enter_context注册上下文管理器到栈中。ExitStack退出时按注册逆序清理所有资源。

ExitStack的回调注册：

stack = ExitStack()
stack.callback(lambda: print("cleanup 1"))
stack.callback(lambda: print("cleanup 2"))
stack.close()
# 输出: cleanup 2 cleanup 1

回调按注册逆序调用。push与callback类似但接受额外的参数。

ExitStack.pop_all转移所有权：

def transaction_context():
stack = ExitStack()
stack.enter_context(start_transaction())
def commit():
if success:
stack.pop_all()
commit_transaction()
stack.close()
return stack, commit

ExitStack的enter_context的异步版本在AsyncExitStack中。

suppress忽略指定异常：

from contextlib import suppress
import os

with suppress(FileNotFoundError):
os.remove('temporary_file.txt')

等效于：

try:
os.remove('temporary_file.txt')
except FileNotFoundError:
pass

suppress的实现接受多个异常类型，在__exit__中捕获并返回True抑制。

redirect_stdout和redirect_stderr捕获输出：

from contextlib import redirect_stdout, redirect_stderr
import io

f = io.StringIO()
with redirect_stdout(f), redirect_stderr(f):
print("this goes to f")
print("this also", file=sys.stderr)
output = f.getvalue()

redirect_stdout替换sys.stdout，redirect_stderr替换sys.stderr。在with块内所有标准输出写入StringIO。

contextlib.ContextDecorator使上下文管理器可作为装饰器：

from contextlib import ContextDecorator

class notify(ContextDecorator):
def __enter__(self):
print('Starting')
return self
def __exit__(self, *exc):
print('Finishing')
return False

@notify()
def my_function():
print('Inside')

my_function()
# Starting
# Inside
# Finishing

ContextDecorator实现了__call__方法，使实例可以作为装饰器使用。decorate方法可以被重写。

contextlib.nullcontext作为不执行任何操作的上下文管理器：

from contextlib import nullcontext

def process(need_db):
ctx = db_session() if need_db else nullcontext()
with ctx as resource:
if resource:
resource.query(...)

nullcontext返回指定的值（默认为None）。在需要条件上下文管理器的场景下替代None检查。

contextlib.chdir临时改变工作目录：

from contextlib import chdir
import os

original = os.getcwd()
with chdir('/tmp'):
print(os.getcwd()) # /tmp
print(os.getcwd()) # 回到原始目录

Python 3.11+。chdir在__enter__时保存pwd，__exit__时恢复。

@contextmanager的异常处理：

@contextmanager
def swallowing():
try:
yield
except ValueError:
print("ValueError caught and swallowed")

with swallowing():
raise ValueError("problem")
print("Continues here")

yield在try块中，异常可以被捕获和处理。如果没有合适的异常处理，异常会在__exit__中重新抛出。

AsyncExitStack用于异步上下文管理器的组合：

from contextlib import AsyncExitStack

async def async_task():
async with AsyncExitStack() as stack:
conn = await stack.enter_async_context(db_connect())
file = await stack.enter_async_context(open_async('file'))
return await process(conn, file)