上下文管理器的核心就是两个魔法方法:enter 和 exit。当你用 with 语句包裹一个对象时,Python 会按约定好的顺序调用这两个方法,从而实现资源获取与释放的自动管理。
基本实现模板
任何一个类只要定义了这两个方法,就可以直接用于 with 语句:
class MyContext:
def __enter__(self):
# 进入 with 块时执行:可以获取资源、打开连接等
print('进入上下文,获取资源')
return self # 返回值会赋给 as 后的变量(如果写了的话)
def __exit__(self, exc_type, exc_val, exc_tb):
# 退出 with 块时必定执行:释放资源、关闭连接等
print('退出上下文,释放资源')
# 返回 False(默认)表示不压制异常,异常会继续向外抛出
# 返回 True 则表示异常已被处理,不再向外传播
return False
使用体验:
with MyContext() as ctx:
print('正在使用上下文资源')
# 即使这里抛出异常,__exit__ 依然会被调用
输出:
进入上下文,获取资源
正在使用上下文资源
退出上下文,释放资源
执行流程详解
- 当解释器执行到
with行时,首先调用对象的enter方法。 enter的返回值会绑定到as后面的变量(如果没写as,返回值会被丢弃)。- 接着执行
with内部的代码块。 - 无论代码块是正常结束、遇到
return还是抛出异常,exit都会被调用。 exit接收三个参数:exc_type:异常的类型(无异常时为None)exc_val:异常实例(无异常时为None)exc_tb:traceback 对象(无异常时为None)- 如果
exit返回True,表示异常已经处理完毕,程序会正常继续执行;如果返回False或None,异常会照常抛出。
实战案例:计时上下文管理器
这是一个非常实用的例子——测量某段代码的执行时间:
import time
class Timer:
def __enter__(self):
self.start = time.perf_counter()
return self # 返回自身,方便在 with 块内获取开始时间等信息
def __exit__(self, exc_type, exc_val, exc_tb):
self.end = time.perf_counter()
self.elapsed = self.end - self.start
print(f'执行耗时:{self.elapsed:.4f} 秒')
return False # 不压制任何异常
用法:
with Timer():
total = sum(range(10_000_000))
输出类似:执行耗时:0.1234 秒。无论 sum 是否出错,你都能看到计时结果(异常下也依然会输出耗时,然后异常再向上抛出)。
实战案例:文件写入自动刷新与关闭
内置 open 已经支持上下文管理,但我们可以自己封装一个,自动在退出时 flush 缓冲区并显式关闭,或者加上额外的日志记录:
class ManagedFile:
def __init__(self, filename, mode='w'):
self.filename = filename
self.mode = mode
def __enter__(self):
self.file = open(self.filename, self.mode)
return self.file
def __exit__(self, exc_type, exc_val, exc_tb):
print('正在关闭文件...')
self.file.close()
return False
使用:
with ManagedFile('test.txt', 'w') as f:
f.write('Hello Python')
实际生产代码里写这样简单封装的情况不多,但这个例子可以帮助你理解 enter 返回的是 open 的文件对象,as f 拿到的就是它。
异常处理测试
来看看 exit 如何处理异常:
class SuppressValueError:
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is ValueError:
print(f'捕获并压制了 ValueError: {exc_val}')
return True # 压制异常
return False # 其他异常正常抛出
测试:
with SuppressValueError():
raise ValueError('一个被压制的异常')
print('程序继续执行') # 会正常打印
with SuppressValueError():
raise TypeError('没有压制的异常') # 会直接抛出 TypeError
输出:
捕获并压制了 ValueError: 一个被压制的异常
程序继续执行
Traceback (most recent call last):
...
TypeError: 没有压制的异常
这里 ValueError 被吞掉了,所以程序继续执行;TypeError 没有被特别处理,于是照常抛出,中断了后续流程。
总结
实现 enter / exit 就是给类赋予“上下文感知”的能力。记住几个关键点:
enter:负责准备资源,返回值交给as变量。exit:负责清理资源,必定执行,可决定是否压制异常。- 一般返回
False:除非你非常清楚自己要吞掉异常,否则不要轻易返回True。
掌握了这两个方法,你就真正理解了 with 语句的底层原理,之后无论是自己写资源管理类,还是阅读框架源码中的上下文管理器,都会游刃有余。