吃透Python上下文管理器

什么是上下文管理器？

我们常见的with open操作文件，就是一个上下文管理器。如：

with open(file, 'rb') as f:
    text = f.read()

那上下文管理器具体的定义是什么呢？

上下文管理器：是指在一段代码执行之前执行一段代码，用于一些预处理工作；执行之后再执行一段代码，用于一些清理工作。

比如刚提到的文件操作，打开文件进行读写，读写完之后需要将文件关闭。很明显用到了上下文管理器。主要依靠__enter__、__exit__这两个”魔术方法”实现。

__enter__(self)

Defines what the context manager should do at the beginning of the block created by the with statement. Note that the return value of enter is bound to the target of the with statement, or the name after the as.

__exit__(self, exception_type, exception_value, traceback)

Defines what the context manager should do after its block has been executed (or terminates). It can be used to handle exceptions, perform cleanup, or do something always done immediately after the action in the block. If the block executes successfully, exception_type, exceptionvalue, and traceback will be None. Otherwise, you can choose to handle the exception or let the user handle it; if you want to handle it, make sure _exit__ returns True after all is said and done. If you don’t want the exception to be handled by the context manager, just let it happen.

当我们需要创建一个上下文管理器类型的时候，就需要实现__enter__和__exit__方法，这对方法称为上下文管理协议（Context Manager Protocol），定义了一种运行时上下文环境。

基本语法

with EXPR as VAR:
    BLOCK

这里就是一个标准的上下文管理器的使用逻辑，其中的运行逻辑：

（1）执行EXPR语句，获取上下文管理器（Context Manager）

（2）调用上下文管理器中的__enter__方法，该方法执行一些预处理工作。

（3）这里的as VAR可以省略，如果不省略，则将__enter__方法的返回值赋值给VAR。

（4）执行代码块BLOCK，这里的VAR可以当做普通变量使用。

（5）最后调用上下文管理器中的的__exit__方法。

（6）__exit__方法有三个参数：exc_type, exc_val, exc_tb。如果代码块BLOCK发生异常并退出，那么分别对应异常的type、value 和 traceback。否则三个参数全为None。

（7）__exit__方法的返回值可以为True或者False。如果为True，那么表示异常被忽视，相当于进行了try-except操作；如果为False，则该异常会被重新raise。

如何自己实现上下文管理器？

简单来说，如果一个类中，实现了__enter__和__exit__方法，那么这个类就是上下文管理器。

class Contextor():
    def __enter__(self):    
        print('程序的预处理开始啦！')
        return self     # 作为as说明符指定的变量的值

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('正在进行收尾！')

    def func(self):
        print('程序进行中....')


with Contextor() as var:
    var.func()
   
# 输出
程序的预处理开始啦！
程序进行中....
正在进行收尾！

从这个示例可以很明显的看出，在编写代码时，可以将资源的连接或者获取放在__enter__中，而将资源的关闭写在__exit__ 中。

为什么要使用上下文管理器？

在我看来，这和 Python 崇尚的优雅风格有关。

1. 可以以一种更加优雅的方式，操作（创建/获取/释放）资源，如文件操作、数据库连接；
2. 可以以一种更加优雅的方式，处理异常；

第二种，会被大多数人所忽略。这里着重介绍下。

在处理异常时，通常都是使用try...except...来进行异常处理的。这就可能会出现在程序的主逻辑中有大量的异常处理代码，这会大大影响程序的可读性。

好的做法可以通过with将异常处理隐藏起来。

我们以1/0举例（1/0必然会抛出错误）

class Resource():
    def __enter__(self):
        print('===connect to resource===')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('===close resource connection===')
        return True

    def operate(self):
        1/0

with Resource() as res:
    res.operate()
    
# 输出结果
===connect to resource===
===close resource connection===

运行发现，并没有出异常。

这就是上下文管理器的强大之处，异常可以在__exit__进行捕获，并自己决定如何处理。在__exit__ 里返回 True（没有return 就默认为 return False），就相当于告诉 Python解释器，这个异常我们已经捕获了，不需要再往外抛了。

在 写__exit__ 函数时，需要注意的事，它必须要有这三个参数：

• exc_type：异常类型
• exc_val：异常值
• exc_tb：异常的错误栈信息

当主逻辑代码没有报异常时，这三个参数将都为None。

如何处理自行处理异常

我们以上面的代码为例，在__exit__加入判断异常的逻辑，如果发生异常，则打印异常信息。

class Resource():
    def __enter__(self):
        print('===connect to resource===')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            print(f'出现异常：{exc_type}：{exc_val}')
        print('===close resource connection===')
        return True

    def operate(self):
        1/0

with Resource() as res:
    res.operate()
    
# 输出
===connect to resource===
出现异常：<class 'ZeroDivisionError'>：division by zero
===close resource connection===

如何更好的使用上下文管理器？

在Python中有一个专门用于实现上下文管理的标准库contextlib。

有了 contextlib 创建上下文管理的最好方式就是使用 contextmanager 装饰器，通过 contextmanager 装饰一个生成器函数，yield 语句前面的部分被认为是__enter__() 方法的代码，后面的部分被认为是 __exit__()方法的代码。

我们以打开文件为例：

import contextlib

@contextlib.contextmanager
def open_func(file_name):
    # __enter__方法
    print('open file:', file_name, 'in __enter__')
    file_handler = open(file_name, 'r')
    
    # 【重点】：yield
    yield file_handler

    # __exit__方法
    print('close file:', file_name, 'in __exit__')
    file_handler.close()
    return

with open_func('/Users/MING/mytest.txt') as file_in:
    for line in file_in:
        print(line)

如果要处理异常，将上面代码改写成下面的样子。

import contextlib

@contextlib.contextmanager
def open_func(file_name):
    # __enter__方法
    print('open file:', file_name, 'in __enter__')
    file_handler = open(file_name, 'r')

    try:
        yield file_handler
    except Exception as exc:
        # deal with exception
        print('the exception was thrown')
    finally:
        print('close file:', file_name, 'in __exit__')
        file_handler.close()

        return

with open_func('test.txt') as file_in:
    for line in file_in:
        1/0
        print(line)

参考：

https://juejin.im/post/5c87b165f265da2dac4589cc
https://www.cnblogs.com/linxiyue/p/3855751.html
https://runnerliu.github.io/2018/01/02/pythoncontextmanager/

扫描下面二维码，关注公众号， 每周定期与您分享原创的、有深度的Python知识点。