Python中必须知道的知识点：上下文管理器

我不是女神ヾ 2022-09-14 13:27 119阅读 0赞

在Python中让自己创建的函数、类、对象支持with语句，就实现了上线文管理协议。我们经常使用with open(file, “a+”) as f:这样的语句，无需手动调用f.close()关闭文件。这种用法不仅优雅，而且避免遗忘释放资源，十分方便。所以，当操作某些资源时，需要对资源的获取与释放进行自动操作，就可以用上线文管理器。比如：数据库的连接，查询，关闭处理；socket的连接和断开。本篇主要介绍，如何让自己创建的类、对象、函数等支持with语句，详细请看下文。

**1 让对象支持上下文管理协议**

在类中，实现`__enter__()`和`__exit__()`方法，类创建的对象就支持with语句。如下：

class A:
        def __enter__(self):
            print("in __enter__")
            return [1, 2, 3]
     
        def __exit__(self, exc_type, exc_val, exc_tb):
            print("in __exit__")
     
    obj = A()
    with obj as o:
        print("看看o是什么：", o)
        print("do something")
    print("end")

注意本例的输出结果的顺序：

in __enter__
    看看o是什么： [1, 2, 3]
    do something
    in __exit__
    end

（1）当执行 with 语句的时候，对象的 `__enter__()`方法被触发， 它返回的值(如果有的话)会被赋值给 as 声明的变量。对应输出应该是`in __enter__`和将\[1,2,3\]赋值给`o字母o`。

（2）然后，with 语句块里面的代码开始执行。对应的输出是`看看o是什么： [1, 2, 3]`和`do someting`。

（3）最后，`__exit__()`方法被触发进行清理工作。对应的输出是`in __exit__`。

补充说明：  
`__exit__()`方法的第三个参数包含了异常类型、异常值和追溯信息(如果有的话)。`__exit__()`方法能自己决定怎样利用这个异常信息，或者忽略它并返回一个None值。

如果 `__exit__()` 返回 True ，那么异常会被清空，就好像什么都没发生一样， with 语句后面的程序继续在正常执行。

上面的例子还不支持多个with嵌套使用，下面是一个可以嵌套使用with语句的例子：

''' 学习中遇到问题没人解答？小编创建了一个Python学习交流QQ群：531509025 寻找有志同道合的小伙伴，互帮互助,群里还有不错的视频学习教程和PDF电子书！ '''
    from socket import socket, AF_INET, SOCK_STREAM
     
    class Connection:
        def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
            self.address = address
            self.family = family
            self.type = type
            self.connections = []
     
        def __enter__(self):
            sock = socket(self.family, self.type)
            sock.connect(self.address)
            self.connections.append(sock)
            return sock
     
        def __exit__(self, exc_ty, exc_val, tb):
            self.connections.pop().close()
     
     
    conn = Connection(('www.python.org', 80))
    with conn as s1:
        print(s1)
        with conn as s2:
            print(s2)

**2 装饰器版上下文管理器**

上面介绍了在类和对象中实现上下文管理协议，其实Python标准库中contextlib包的@contextmanager装饰器能够轻松实现一个上下文管理器，下例是用其实现统计代码块耗时的上下文管理器：

import time
    from contextlib import contextmanager
    # 来看一个装饰器版本的上下文管理器
    # 检查代码消耗时间块
    @contextmanager
    def timecount(name):
        start = time.time()
        try:
            yield
        finally:
            end = time.time()
            print('{}: {}'.format(name, end - start))
     
    with timecount('cost time:'):
        time.sleep(2)
    # cost time:: 2.000200033187866

timecount（）中，yield之前的代码相当于`__enter__()`方法；yield 之后的代码相当于 `__exit__()`方法。如果有异常会自动在yield一行抛出。

上下文管理器可以应用在事务中：

''' 学习中遇到问题没人解答？小编创建了一个Python学习交流QQ群：531509025 寻找有志同道合的小伙伴，互帮互助,群里还有不错的视频学习教程和PDF电子书！ '''
    # 更高级的事务管理
    @contextmanager
    def list_transaction(orig_list):
        working = list(orig_list)
        yield working
        orig_list[:] = working
     
    lis = [1, 2, 3]
    with list_transaction(lis) as work:
        work.append(5)
        work.append(6)
     
    print(lis)
    # [1, 2, 3, 5, 6]

一旦with语句内有异常产生，之前的操作不会生效：

lis = [1, 2, 3]
    with list_transaction(lis) as work:
        work.append(5)
        work.append(6)
        print(lis)
        raise RuntimeError("回滚")

如果在交互式命令行中打印lis,依然会发现lis的内容没有改变，也就是说，with语句中的代码出现异常，with语句中的操作都不会生效，只有with无异常退出，才会生效。对于事务管理来说是比较有用的。

**3 总结**

（1）当操作一些需要打开、连接、断开或释放的资源时，让对象或函数支持with语句十分方便、省事、优雅。如数据库的链接断开、套接字的连接断开、事务、锁资源等。

（2）类中实现`__enter__()`和`__exit__()`方法，即可实现上下文管理协议，对象可使用with语句。

（3）通过contextlib包中的@contextmanager装饰器装饰一个函数，该函数即可使用with语句。