with 语句
在我们日常使用场景中,经常会操作一些资源,比如文件对象、数据库连接、Socket 连接等。资源操作完了之后,不管操作的成功与否,最重要的事情就是关闭该资源,否则资源打开太多而没有关闭,程序会报错。以文件操作为例,通常我们会这样写:
f = open('file.txt', 'w', encoding='utf-8')
try:
f.write("Hello")
finally:
f.close()
但既然 close 方法是必须的操作,那就没必要显式地调用,所以 Python 给我们提供了一种更优雅的方式,使用 with 语句:
withopen('file.txt', 'w', encoding='utf-8') as f:
f.write("Hello")
在退出 with 语句下的代码块之后,f 对象会自动执行自己的 close 方法,实现资源的释放,简洁优雅。
这里比原文多写了一个 encoding='utf-8'。文本模式下不传 encoding 会使用与平台相关的默认编码(Windows 上往往不是 UTF-8),Python 3.10 起加入了 EncodingWarning 来提示这件事,用 python -X warn_default_encoding 运行就能看到警告。写新代码时显式指定编码是更稳妥的习惯。
上下文管理器原理
上下文管理器实际是内部实现了 __enter__ 和 __exit__ 方法的对象。
当我们使用 with 语法时:
__enter__() 方法:返回一个值,可以将它赋值给 as 后面的对象,例如上面的 f;
__exit__() 方法:with 语句退出或者发生异常时会执行这个方法。
1、__enter__ 方法说明
上下文管理器的 __enter__ 方法是可以带返回值的,默认返回 None,这个返回值通过 with ... as ... 中的 as 赋给它后面的那个变量,所以 with EXPR as VAR 就是将 EXPR 对象 __enter__ 方法的返回值赋给 VAR。
当然 with...as... 并非固定组合,单独使用 with... 也是可以的,上下文管理器的 __enter__ 方法还是正常执行,只是这个返回值并没有赋给一个变量,with 下面的代码块也不能使用这个返回值。
如果对象没有实现这两个方法就拿去 with,3.11 之后的报错信息比过去清楚很多,会直接告诉你缺了哪个方法:
classOnlyEnter:
def__enter__(self):
returnself
with OnlyEnter():
pass
# TypeError: 'OnlyEnter' object does not support the context manager protocol (missed __exit__ method)
2、__exit__ 方法说明
上下文管理器的 __exit__ 方法接收 3 个参数 exc_type、exc_val、exc_tb,如果代码块 BLOCK 发生了异常 e 并退出,这 3 个参数分别为 type(e)、e 本身(异常实例)、e.__traceback__,否则都为 None。
原文这里说第二个参数是 str(e),这个说法是错的。实际传入的是异常实例,不是字符串,实测可以验证:
classTrack:
def__enter__(self):
returnself
def__exit__(self, exc_type, exc_val, exc_tb):
print(exc_type, repr(exc_val), type(exc_val), exc_tb isnotNone)
returnFalse
with Track():
raise ValueError("boom")
# <class 'ValueError'> ValueError('boom') <class 'ValueError'> True
# 之后异常继续向外抛出
这个区别在实际写代码时很关键,因为你可以直接对 exc_val 做 isinstance 判断、读它的属性、或者重新 raise,如果它只是个字符串这些都做不到。
同样 __exit__ 方法也是可以带返回值的,这个返回值应该是一个布尔类型 True 或 False,默认为 None(即 False)。如果为 False,异常会被抛出,用户需要进行异常处理。如果为 True,则表示忽略该异常。
classSwallow:
def__enter__(self):
returnself
def__exit__(self, exc_type, exc_val, exc_tb):
print("swallow", exc_type)
returnTrue
with Swallow():
raise KeyError("k")
print("after swallow")
# swallow <class 'KeyError'>
# after swallow
一个上下文管理器一般使用如下:
with EXPR as VAR:
BLOCK
上述代码的执行过程等价于:
import sys
manager = EXPR
enter = type(manager).__enter__
exit = type(manager).__exit__
value = enter(manager)
hit_except = False
try:
VAR = value
BLOCK
except BaseException:
hit_except = True
ifnot exit(manager, *sys.exc_info()):
raise
finally:
ifnot hit_except:
exit(manager, None, None, None)
原文给的等价代码是 try/finally 加一个无参的 ContextManager.__exit__(),这只是个便于理解的粗略示意,实际上有三处不成立,写自定义管理器时会踩到:
__exit__ 必须收到 3 个参数。文件对象的 __exit__ 签名是 (*args),所以 f.__exit__() 空参调用碰巧能跑;但自己定义的 def __exit__(self, exc_type, exc_val, exc_tb) 空参调用会直接 TypeError。
无参调用等于永远告诉管理器「没有异常」,__exit__ 就丧失了处理异常和吞掉异常的能力,返回 True 也起不了作用。
特殊方法是从类型上查找的(type(manager).__enter__),不是从实例的 __dict__,所以给实例挂一个 __enter__ 属性对 with 无效。
f 对象就是把自己的 close 方法定义在了它的 __exit__ 方法内部,实现了代码块 BLOCK 执行完之后自动关闭自身。
自定义上下文管理器
下面我们定义一个文件类,内部实现了 __enter__ 和 __exit__ 两个方法:
classFile:
def__init__(self, filename, mode, encoding='utf-8'):
self.filename = filename
self.mode = mode
self.encoding = encoding
self.f = None
def__enter__(self):
print("进入")
self.f = open(self.filename, self.mode, encoding=self.encoding)
returnself.f
def__exit__(self, exc_type, exc_val, exc_tb):
print("退出")
self.f.close()
returnFalse
这时候 File 类就是一个上下文管理器。
原文把 __exit__ 写成 def __exit__(self, exc_type=None, exc_val=None, exc_tbs=None),给三个参数都加了默认值,目的是让后面 __exit__() 的空参调用不报错。真实的协议调用总会传满 3 个参数,所以默认值没有必要,反而掩盖了「无参调用是错的」这个事实,这里改回标准签名。另外 return False 显式写出来,表示不吞异常(不写也一样,None 即假值)。
我们分别通过 with 语句和 try/finally 语句使用 File 类对文件进行写入操作。
通过 with 语句执行:
with File('file.txt', 'w') as f:
print("正在写入...")
f.write('Hello')
控制台输出:
进入
正在写入...
退出
并得到了一个写了 Hello 的 file.txt 文件。
通过 try/finally 语句执行:
manager = File('file.txt', 'w')
var = manager.__enter__()
try:
print("正在写入...")
var.write('Hello')
finally:
manager.__exit__(None, None, None)
控制台输出:
进入
正在写入...
退出
并得到了一个写了 Hello 的 file.txt 文件。
两者输出一致,所以验证了上一节中执行过程的等价关系是正确的。注意这里的 __exit__ 必须补上三个 None,这也是上一节说的那处改动。
如果希望自定义的管理器带上一点标准行为,可以继承 contextlib.AbstractContextManager,它提供了默认的 __enter__(返回 self)和 __exit__(返回 None),只需要覆盖需要的那一个:
from contextlib import AbstractContextManager
classTimer(AbstractContextManager):
def__exit__(self, exc_type, exc_val, exc_tb):
print("退出")
returnFalse
with Timer() as t:
print("t 就是实例本身:", t)
contextmanager 装饰器
Python 还提供了一个 contextmanager 装饰器,允许用户将一个生成器定义为上下文管理器。该装饰器将生成器中的代码通过 yield 语句分成两部分,yield 之前的代码为 __enter__ 方法,yield 之后的代码为 __exit__ 方法,yield 的返回值即 __enter__ 方法的返回值,用于赋给 as 后的变量。
下面我们通过 contextmanager 装饰器也实现一个关于文件的上下文管理器:
from contextlib import contextmanager
@contextmanager
defopen_file(filename, mode, encoding='utf-8'):
print('进入')
f = open(filename, mode, encoding=encoding)
try:
yield f
finally:
print('退出')
f.close()
说明:这里使用 try/finally 是确保 yield 的过程中就算出现异常,文件也能正常关闭。当然这里也能处理异常,使用 try/except/finally 即可。
通过 with 语句执行:
with open_file('file.txt', 'w') as f:
print("正在写入...")
f.write('Hello')
执行结果跟之前的上下文管理器执行结果一致,说明 contextmanager 装饰器也能定义一个上下文管理器。
有两个细节值得补充。一是生成器版本的管理器是一次性的,调用一次函数得到的对象只能进一个 with;想复用必须每次重新调用函数:
cm = open_file('file.txt', 'w')
with cm:
pass
with cm: # 再次进入会报错
pass
# AttributeError: '_GeneratorContextManager' object has no attribute 'args'
二是生成器的形状必须正好 yield 一次,否则 with 会抛 RuntimeError,报错信息很直白:一次都不 yield 是 generator didn't yield,yield 两次是 generator didn't stop。
@contextmanager 产生的对象同时也是装饰器(继承了 ContextDecorator),可以直接拿去装函数,省掉在函数体里再套一层 with:
from contextlib import contextmanager
@contextmanager
deftag():
print("<t>")
yield
print("</t>")
@tag()
defhello():
print("hi")
hello()
# <t>
# hi
# </t>
contextlib 里其他好用的工具
标准库里已经有一批现成的上下文管理器,日常写代码比自己手写 __enter__/__exit__ 更省事:
from contextlib import suppress, nullcontext, closing, ExitStack, chdir
import os
with suppress(ZeroDivisionError): # 吞掉指定异常,替代空的 try/except/pass
1 / 0
with nullcontext(42) as v: # 3.7+,什么都不做的占位管理器
print(v) # 42
with ExitStack() as stack: # 数量不固定的资源,统一注册、逆序释放
files = [stack.enter_context(closing(obj)) for obj in objs]
with chdir('C:/Windows'): # 3.11+,临时切换工作目录并自动切回
print(os.getcwd())
同时管理多个资源时,从 Python 3.10 起可以用括号把它们分行写,不再需要反斜杠续行:
with (
open('a.txt', encoding='utf-8') as fa,
open('b.txt', 'w', encoding='utf-8') as fb,
):
fb.write(fa.read())