Python contextlib模块高级使用技巧contextlib提供了比with语句更丰富的上下文管理工具。contextmanager装饰器将一个生成器函数转为上下文管理器from contextlib import contextmanagercontextmanagerdef managed_resource(*args, **kwargs):resource acquire_resource(*args, **kwargs)try:yield resourcefinally: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, kwdsdef __enter__(self):try:return next(self.gen)except StopIteration:raise RuntimeError(generator didnt yield)def __exit__(self, type, value, traceback):if type is None:try:next(self.gen)except StopIteration:return Falseelse:raise RuntimeError(generator didnt stop)else:if value is None:value type()try:self.gen.throw(type, value, traceback)except StopIteration as exc:return exc is not valueexcept RuntimeError as exc:if exc is value:return Falseif exc.__cause__ is value:return Falseraiseexcept BaseException:if sys.exc_info()[1] is value:return Falseraisereturn True__enter__调用next启动生成器到yield。__exit__恢复生成器执行清理代码。异常通过throw传递到生成器内部。ExitStack组合管理多个上下文管理器from contextlib import ExitStackdef 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()raiseenter_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, commitExitStack的enter_context的异步版本在AsyncExitStack中。suppress忽略指定异常from contextlib import suppressimport oswith suppress(FileNotFoundError):os.remove(temporary_file.txt)等效于try:os.remove(temporary_file.txt)except FileNotFoundError:passsuppress的实现接受多个异常类型在__exit__中捕获并返回True抑制。redirect_stdout和redirect_stderr捕获输出from contextlib import redirect_stdout, redirect_stderrimport iof io.StringIO()with redirect_stdout(f), redirect_stderr(f):print(this goes to f)print(this also, filesys.stderr)output f.getvalue()redirect_stdout替换sys.stdoutredirect_stderr替换sys.stderr。在with块内所有标准输出写入StringIO。contextlib.ContextDecorator使上下文管理器可作为装饰器from contextlib import ContextDecoratorclass notify(ContextDecorator):def __enter__(self):print(Starting)return selfdef __exit__(self, *exc):print(Finishing)return Falsenotify()def my_function():print(Inside)my_function()# Starting# Inside# FinishingContextDecorator实现了__call__方法使实例可以作为装饰器使用。decorate方法可以被重写。contextlib.nullcontext作为不执行任何操作的上下文管理器from contextlib import nullcontextdef 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 chdirimport osoriginal os.getcwd()with chdir(/tmp):print(os.getcwd()) # /tmpprint(os.getcwd()) # 回到原始目录Python 3.11。chdir在__enter__时保存pwd__exit__时恢复。contextmanager的异常处理contextmanagerdef swallowing():try:yieldexcept ValueError:print(ValueError caught and swallowed)with swallowing():raise ValueError(problem)print(Continues here)yield在try块中异常可以被捕获和处理。如果没有合适的异常处理异常会在__exit__中重新抛出。AsyncExitStack用于异步上下文管理器的组合from contextlib import AsyncExitStackasync 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)