Python 上下文管理器

with open(file='log.txt', mode='r', encoding='gbk') as f:
    r = f.read(5)
    print(r)

在打开文件或连接数据库时，我们往往会使用with语句。with可以帮我们自动处理异常、关闭连接等。那么with是怎么实现这些功能的呢？

1.上下文管理协议

首先，不是所有对象都可以跟在with后面，下面例子我们随便定义一个类并使用with。

class A:
    def __init__(self):
        pass

    def fun(self):
        print('随便')


with A() as a:
    a.fun()


# 输出：AttributeError: __enter__

通过报错可以看到，缺少__enter__方法。实际上，要实现一个对象能被with调用，需要至少实现__enter__和__exit__方法。

# 上下文管理器协议
class MyContext:
    def __enter__(self):
        print("enter")
        # 这个返回的self会被作为as的对象
        return self

    # 参数是默认写法，可以捕获一些异常
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("exit")

    @staticmethod
    def do_something():
        print("doing something")


# 定义一个具有打开（__enter__）和释放（__exit__）资源的类，使用with语句时会自动调用
with MyContext() as ctx:
    ctx.do_something()


# 输出：enter
#       doing something
#       exit

2.使用细节

__enter__方法会被首先调用，其返回的对象被赋给as后的变量，比如我们可以直接返回一个基本类型

class MyContext:
    def __enter__(self):
        print("enter")
        # 这个返回的self会被作为as的对象
        return '返回一个字符串'

    # 参数是默认写法，可以捕获一些异常
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("exit")

    @staticmethod
    def do_something():
        print("doing something")


with MyContext() as ctx:
    print(ctx)


# 输出：enter
#       返回一个字符串
#       exit

而我们一开始提到with可以处理异常，其实现就是靠__exit__捕获，我们可以制造一个异常输出看看。

class MyContext:
    def __enter__(self):
        print("建立连接")
        # 这个返回的self会被作为as的对象
        return self

    # 参数是默认写法，可以捕获一些异常
    def __exit__(self, exc_type, exc_val, exc_tb):
        print('错误信息', exc_type, exc_val, exc_tb)
        if type(exc_val) == ZeroDivisionError:
            print("遇到错误，安全关闭连接")

    @staticmethod
    def do_something():
        print("doing something")


with MyContext() as ctx:
    1/0
    ctx.do_something()


# 输出：建立连接
#       错误信息 <class 'ZeroDivisionError'> division by zero <traceback object at # 
        0x0000019A5B3EC248>
#       遇到错误，安全关闭连接

3.使用装饰器实现上下文管理

上面的例子可以看到，enter和exit都是不是对象的核心方法，更像是一层装饰。这一点就和装饰器（Python 装饰器_Jiangugu的博客-CSDN博客）的特点很吻合了。同时，通过上下文管理器执行顺序可以看到，先执行了对象的enter方法->然后执行with下面的语句体->然后才是exit方法。这种执行代码段的切换，就和生成器（Python 异步编程之——协程_Jiangugu的博客-CSDN博客）特点符合了。因此，基于装饰器和生成器就可以实现上下文管理。

from contextlib import contextmanager


# 使用装饰器使函数具有上下文管理功能
@contextmanager
def file_open(file):
    print("{} 打开".format(file))
    yield ''
    print("{} 关闭".format(file))


with file_open("日志.txt") as f:
    print("读取文件中")


# 输出：日志.txt 打开
#       读取文件中
#       日志.txt 关闭