Python 上下文管理器与 with 语句

with 语句到底保证了什么，怎样编写 __enter__ 和 __exit__，contextlib 提供的简便写法，以及如何有意吞掉异常。

with 存在的意义，就是不管代码块里发生什么，清理工作都会执行。提前返回、抛出异常，或者正常结束，退出代码照样会运行。

任何带有 __enter__ 和 __exit__ 的对象都能这样用，所以你也可以把自己的资源放到同一种语句后面管理。

两个方法各做什么

__enter__ 最先运行，它的返回值就是 as 绑定的对象。__exit__ 在最后运行，会收到当前正在传播的异常（如果有的话）。

# what with actually guarantees
import io

class Tracked:
    def __init__(self):
        self.events = []
    def __enter__(self):
        self.events.append('enter')
        return self
    def __exit__(self, exc_type, exc, tb):
        self.events.append(f'exit ({exc_type.__name__ if exc_type else "no error"})')
        return False

t = Tracked()
with t as handle:
    handle.events.append('body')
print(t.events)

输出：

['enter', 'body', 'exit (no error)']

出错时，清理照样执行

__exit__ 运行完之后，异常继续向外传播。

# cleanup runs even when the body raises
t2 = Tracked()
try:
    with t2:
        t2.events.append('body')
        raise ValueError('boom')
except ValueError as err:
    t2.events.append(f'caught {err}')
print(t2.events)

输出：

['enter', 'body', 'exit (ValueError)', 'caught boom']

关键就在这里。try/finally 也能做到同样的事，但 with 把清理代码放在资源旁边，而不是散落在每个调用处。

返回 True 会吞掉异常

__exit__ 返回真值，意思是“我已经处理了”。

# returning True from __exit__ swallows the exception
class Swallow:
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc, tb):
        return True

with Swallow():
    raise RuntimeError('this never escapes')
print('still running')

输出：

still running

这种用法要少用，而且要想清楚再用。悄悄吞掉所有异常，错误就是这样消失不见的。

contextlib 的简便写法

只含一个 yield 的生成器可以变成上下文管理器。yield 之前的都是准备工作，之后的都是清理工作。

# contextlib turns a generator into a context manager
import contextlib

@contextlib.contextmanager
def tag(name, out):
    out.append(f'<{name}>')
    try:
        yield out
    finally:
        out.append(f'</{name}>')

lines = []
with tag('p', lines):
    lines.append('hello')
print(''.join(lines))

输出：

<p>hello</p>

正是包在 yield 外面的 try/finally，保证了代码块抛异常时清理照样执行。没有它，异常会跳过生成器剩下的部分。

用 suppress 代替 try/except/pass

contextlib.suppress 一眼就能看出意图，而且只占一行。

# contextlib.suppress instead of try/except/pass
with contextlib.suppress(ZeroDivisionError):
    1 / 0
print('suppressed')

输出：

suppressed

同时使用多个管理器

用逗号分隔时，它们从左到右嵌套，按相反顺序关闭。

# several managers on one line
buf1, buf2 = io.StringIO(), io.StringIO()
with buf1 as a, buf2 as b:
    a.write('one')
    b.write('two')
    print(a.getvalue(), b.getvalue())     # read before the block closes them
print('both closed:', buf1.closed, buf2.closed)

输出：

one two
both closed: True True

注意第二行。代码块一结束，文件就关闭了，所以要读的内容得在块内读完。

要点

• __enter__ 返回 as 绑定的对象，__exit__ 一定会运行。
• __exit__ 返回真值会吞掉异常。
• @contextlib.contextmanager 加上 try/finally，是写上下文管理器的简便方法。
• 写在同一行 with 里的管理器按相反顺序关闭。

锁、数据库事务、临时目录、计时器和模拟的配置，在标准库里都是上下文管理器。如果你的代码有准备步骤和对应的收尾步骤，它也应该写成上下文管理器。