上下文管理器

在 Python 开发中，我们经常会使用到 with 语法块，例如在读写文件时，保证文件描述符的正确关闭，避免资源泄露问题。

你有没有思考过， with 背后是如何实现的？我们常常听到的上下文管理器究竟是什么？

这篇文章我们就来学习一下 Python 上下文管理器，以及 with 的运行原理。

with语法块

在讲解 with 语法之前，我们先来看一下不使用 with 的代码如何写？

我们在操作一个文件时，代码可以这么写：

# 打开文件
f = open('file.txt')
for line in f:
    # 读取文件内容 执行其他操作
    # do_something...
# 关闭文件
f.close()

这个例子非常简单，就是打开一个文件，然后读取文件中的内容，最后关闭文件释放资源。

但是，代码这么写会有一个问题：在打开文件后，如果要对读取到的内容进行其他操作，如果操作期间发生了异常，这就会导致文件句柄无法被释放，进而导致资源的泄露。

如何解决这个问题？

也很简单，我们使用 try ... finally 来优化代码：

# 打开文件
f = open('file.txt')
try:
    for line in f:
        # 读取文件内容 执行其他操作
        # do_something...
finally:
    # 保证关闭文件
    f.close()

这么写的好处是，在读取文件内容和操作期间，无论是否发生异常，都可以保证最后能释放文件资源。

但这么优化，代码结构会变得很繁琐，每次都要给代码逻辑增加 try ... finally 才可以，可读性变得很差。

针对这种情况，我们就可以使用 with 语法块来解决这个问题：

with open('file.txt') as f:
    for line in f:
        # do_something...

使用 with 语法块可以完成之前相同的功能，而且这么写的好处是，代码结构变得非常清晰，可读性也很好。

明白了 with 的作用，那么 with 究竟是如何运行的呢？

上下文管理器

首先，我们来看一下 with 的语法格式：

with context_expression [as target(s)]:
    with-body

with 语法非常简单，我们只需要 with 一个表达式，然后就可以执行自定义的业务逻辑。

但是，with 后面的表达式是可以任意写的吗？

答案是否定的。要想使用 with 语法块，with 后面的的对象需要实现「上下文管理器协议」。

什么是「上下文管理器协议」？

一个类在 Python 中，只要实现以下方法，就实现了「上下文管理器协议」：

• __enter__：在进入 with 语法块之前调用，返回值会赋值给 with 的 target
• __exit__：在退出 with 语法块时调用，一般用作异常处理

我们来看实现了这 2 个方法的例子：

class TestContext:

    def __enter__(self):
        print('__enter__')
        return 1

    def __exit__(self, exc_type, exc_value, exc_tb):
        print('exc_type: %s' % exc_type)
        print('exc_value: %s' % exc_value)
        print('exc_tb: %s' % exc_tb)

with TestContext() as t:
    print('t: %s' % t)
    
# Output:
# __enter__
# t: 1
# exc_type: None
# exc_value: None
# exc_tb: None

在这个例子中，我们定义了 TestContext 类，它分别实现了 __enter__ 和 __exit__ 方法。

这样依赖，我们就可以把 TestContext 当做一个「上下文管理器」来使用，也就是通过 with TestContext() as t 方式来执行。

从输出结果我们可以看到，具体的执行流程如下：

• __enter__ 在进入 with 语句块之前被调用，这个方法的返回值赋给了 with 后的 t 变量
• __exit__ 在执行完 with 语句块之后被调用

如果在 with 语句块内发生了异常，那么 __exit__ 方法可以拿到关于异常的详细信息：

• exc_type：异常类型
• exc_value：异常对象
• exc_tb：异常堆栈信息

我们来看一个发生异常的例子，观察 __exit__ 方法拿到的异常信息是怎样的：

with TestContext() as t:
    # 这里会发生异常
    a = 1 / 0 
    print('t: %s' % t)

# Output:
# __enter__
# exc_type: <type 'exceptions.ZeroDivisionError'>
# exc_value: integer division or modulo by zero
# exc_tb: <traceback object at 0x10d66dd88>
# Traceback (most recent call last):
#   File "base.py", line 16, in <module>
#     a = 1 / 0
# ZeroDivisionError: integer division or modulo by zero

从输出结果我们可以看到，当 with 语法块内发生异常后，__exit__ 输出了这个异常的详细信息，其中包括异常类型、异常对象、异常堆栈。

如果我们需要对异常做特殊处理，就可以在这个方法中实现自定义逻辑。

回到最开始我们讲的，使用 with 读取文件的例子。之所以 with 能够自动关闭文件资源，就是因为内置的文件对象实现了「上下文管理器协议」，这个文件对象的 __enter__ 方法返回了文件句柄，并且在 __exit__ 中实现了文件资源的关闭，另外，当 with 语法块内有异常发生时，会抛出异常给调用者。

伪代码可以这么写：

class File:

    def __enter__(self):
        return file_obj

    def __exit__(self, exc_type, exc_value, exc_tb):
        # with 退出时释放文件资源
        file_obj.close()
        # 如果 with 内有异常发生 抛出异常
        if exc_type is not None:
            raise exception

这里我们小结一下，通过对 with 的学习，我们了解到，with 非常适合用需要对于上下文处理的场景，例如操作文件、Socket，这些场景都需要在执行完业务逻辑后，释放资源。