理解 Python 的上下文管理器

任何 Python 教程,必然会讲解如何打开一个文件。而任何提到打开文件的地方,都必然会推荐用 with 来操作文件的读写。比如说这里有一篇非常优秀的教程,文中提到

在 Python 中,文件读写是通过 open() 函数打开的文件对象完成的。使用 with 语句操作文件 IO 是个好习惯。

并且给出了详细的代码示例。但为什么 with 关键字能在结束这个 block 的时候自动调用 close() 呢?让我们去一探究竟。

首先来看一下 with 是怎么来的。 with 被提出是在 PEP 343 ,其中有段对于 with 操作的详细说明。简单来说,执行下面代码的前提是,要求 EXPR 的类实现了 __enter____exit__ 方法。

with EXPR as VAR:
    BLOCK

进入 with block 之后,第一件事就是把 __enter__ 的返回值赋给 VAR ,然后执行 BLOCK 的内容,无论能否顺利执行,最终都会执行 __exit__ 方法来“收拾残局”,在我们的情况下,即关闭文件。

了解了 with 的原理,那么 open 函数又是如何提供 __enter____exit__ 方法的呢?让我们去源码里找找蛛丝马迹。当我们写下 open() 某个文件时, Python 实际上调用的是这里,这个函数实质上返回的是类 TextIOWrapper 的实例。我们可以通过 Python shell 来证实。

>>> open("test.txt")
<_io.TextIOWrapper name='test.txt' mode='r' encoding='UTF-8'>

这个类继承自 TextIOBase ,而 TextIOBase 又继承自 IOBase ,这个类就是所有文件 IO 类的基类,所有文件读写的类都继承了它。通过阅读这个基类,我们发现其中有这么一段代码

...
    ### Context manager ###

    def __enter__(self):  # That's a forward reference
        """Context management protocol.  Returns self (an instance of IOBase)."""
        self._checkClosed()
        return self

    def __exit__(self, *args):
        """Context management protocol.  Calls close()"""
        self.close()
...

于是我们找到了原因,当通过 with 来打开一个文件时,我们得到的是一个 IOBase 子类的实例,这个实例提供各种读写文件的方法,当退出 with 代码块时,调用文件关闭方法。这样我们就不用编写大段的 try...except...finally 来确保文件安全关闭了。

PEP 343 说了,只要一个类能提供 __enter____exit__ 方法,我们就能用 with 来保证某个操作在代码执行完毕后能继续执行,作为收尾。我们试着写一个简单的类,来证明我们的理解是对的。要求在 with 结束后,打印出一句我最喜欢的诗句。

>>> class Foo(object):
...     def __init__(self, bar):
...         self.bar = bar
...     def __enter__(self):
...         return self
...     def __exit__(self, *args):
...         print("苟利国家生死以,岂因祸福避趋之")
...
>>> with Foo("naive") as foo:
...     print(foo.bar)
...
naive
苟利国家生死以,岂因祸福避趋之

非常简单的一个类,就保证了每次代码执行之后,都能在屏幕上打印出这样一句话。但这样仍然稍显麻烦,毕竟我想的是能少些两句就少些两句,每次专门写这样一个类,还必须实现那两个方法,想想就挺啰嗦,有没有更好的办法?办法就在这句“上下文管理器”的注释里 ### Context manager ### ,稍微搜索一下,我们就能找到这样一篇文档,参照代码示例,我们不难写出一段简单的代码来实现同样的功能。

>>> from contextlib import contextmanager
>>>
>>> @contextmanager
... def foo(bar):
...     yield bar
...     print("苟利国家生死以,岂因祸福避趋之")
...
>>> with foo("naive") as f:
...     print(f)
...
naive
苟利国家生死以,岂因祸福避趋之

这里 yield 抛出的是我们真正感兴趣的内容,后续则是 with 块结束后我们必须进行的操作。