Python核心技术与实战——二一|巧用上下文管理器和with语句精简代码

我们在Python中对于with的语句应该是不陌生的，特别是在文件的输入输出操作中，那在具体的使用过程中，是有什么引伸的含义呢？与之密切相关的上下文管理器(context manager)又是什么呢？

什么是上下文管理器

在任何一种编程语言里，文件的输入输出、数据库的建立连接和断开等操作，都是很常见的资源管理操作。但是资源是有限的，在写程序的时候，我们必须保证这些资源在使用后得到释放，不然就容易造成资源泄漏，轻者系统处理缓慢，重则系统崩溃。

我们看一个例子：

for i in range(100000000):
    f = open('test.txt','w')
    f.write('hello')

我们在循环里打开了100000000个文件，但是在使用完毕后没有进行关闭操作，一运行代码，就报错了。

这就是一个典型的资源泄漏的案例，因为程序中同时打开了太多的文件，占用了太多的资源，造成崩溃。

为了解决这个问题，不同的编程语言都引入了不同的机制，在Python中，对应的解决方法就是上下文管理器(context manager)。上下文管理器能够自动分配资源并释放资源，其中最典型的应用就是with语句，所以上面的代码应该用这种方式来写

for x in range(100000000):
    with open('test.txt','w') as f:
        f.write('hello world')

这样，我们每次打开文件“test.txt"并写入字符以后，这个文件就会自动关闭，相应的资源也就得以释放，可以防止资源泄漏。当然，with的语句也可以用下面的方式来表示

f = open('test.txt','w')
try:
    f.write('hello world')
finally:
    f.close()

这里一定要注意finally的程序段，哪怕在写入的时候发生了异常，他也可以保证文件最终被关闭。不过于with想比较，就显得比较冗余了，并且还容易忽略finally，所以我们平时更倾向于使用with语句。

另外一种很典型的例子，就是Python中的线程锁(threrading.lock类)，比如我们想要获得一个锁，执行相应的操作以后再将其释放，那么代码就应该是这样的

import threading
some_lock = threading.Lock()
some_lock.acquire()
try:
    pass
finally:
    some_lock.release()

而与其对应的with语句就非常简洁了

import threading
some_lock = threading.Lock()
with some_lock:
    pass

从上面两个例子可以发现，使用with语句，可以大大的简化代码结构，有效的避免资源泄漏的发生。

---

上下文管理器的实现

基于类的上下文管理器 

了接了上下文管理的概念和优点以后，我们就通过下面的例子，看看上下文管理器的原理，高清他的内部实现。我们在这里定义一个上下文管理类FileManager，来模拟Python的打开、关闭文件的操作

class FileManager():
    def __init__(self,name,mode):
        print('call __init__ method')
        self.name = name
        self.mode = mode
        self.file = None

    def __enter__(self):
        print('calling __enter__ method')
        self.file = open(self.name,self.mode)
        return self.file
    def __exit__(self,exc_type,exc_val,exc_tb):
        print('call __exit__ method')
        if self.file:
            self.file.close()

with FileManager('test.txt','w') as f:
    print('ready to write to file')
    f.write('hello world')

##########输出##########
call __init__ method
calling __enter__ method
ready to write to file
call __exit__ method

特别注意：当我们用类来创建上下文管理器的时候，必须保证这个类包括下面两个方法：

__enter__()
__exit__()

并且enter方法还要返回需要被管理的资源，方法exit里通常会存在一些释放、清理资源的操作，比如上面这段代码里的关闭文件等。

而当我们用with语句来执行上面这个上下文管理器的时候，会发生下面四个步骤：

1.构造方法__init__()会被调用，程序初始化对象FileManager，使得文件名和操作方式被传入。

2.方法__enter__()被调用，文件被以写入的模式打开，并且返回FileManager对戏那个赋值给变量f

3.字符串被写入文件

4.方法__exit__()被调用，关闭之前打开的文件流。

所以就有了上面列出的输出结果。

另外我们可以看到exit函数里传递了几个参数——exc_type.exc_val,exc_tb，分别表示exception_type,exception_value和traceback。当我们执行含有上下文管理器的with语句的时候，如果有异常抛出，异常的信息就会被包含在上面三个参数中，传给__exit__()函数。

因此，如果我们需要处理一些异常，可以在__exit__()函数中添加相应的代码

    def __exit__(self,exc_type,exc_val,exc_tb):
        print('call __exit__ method')
        if exc_type:
            print(f'exc_type:{exc_type}')
            print(f'exc_value:{exc_val}')
            print(f'exc_traceback:{exc_tb}')
            print('exception handled')
            return True
        if self.file:
            self.file.close()

with FileManager('test.txt','w') as f:
    raise Exception('exception raised').with_traceback(None)

在修改了exit()方法以后我们在with语句中用raise手动抛出异常，我们可以看到代码有下面的输出

call __init__ method
calling __enter__ method
call __exit__ method
exc_type:<class 'Exception'>
exc_value:exception raised
exc_traceback:<traceback object at 0x000001AA82C8FF48>
exception handled

要注意的是，如果exit函数如果没有返回True，呢么异常仍然会被抛出。如果我们确定了异常已经被处理，那么在exit最后要加上True的返回值。

---

同样的，在数据库的连接等操作上，也常常使用上下文管理器来表示，下面是个简化的代码

class DBConnectionManager():
    def __init__(self,hostname,port):
        self.hostname = hostname
        self.port = port
        self.connection = None

    def  __enter__(self):
        self.connection = DBClient(self.hostname,self.port)
        return self

    def __exit__(self,exc_type,exc_val,exc_tb):
        self.connection.close()

with DBConnectionManager('localhost','') as db_client:
    pass

代码的具体含义和上面的例子类似，就不再详细说明了。只要我们写完了DBConnectionManager这个类以后，在每次建立数据库连接时，只要简单的利用with语句就可以了，并不要关系数据库的关闭、异常等，大大的提高了开发效率。

基于生成器的上下文管理器

上面那种基于类的上下文管理器在Python中利用非常广泛，我们在很多项目中都可以看得到，不过Python中的上下文管理器不仅仅局限于此，畜类基于类，它还可以基于生成器实现，我们看看下面的例子：

我们用一个装饰器contextlib.contextmanager来定义自己所需要的基于生成器的上下文管理器，用以支持with语句。同样我们用前面的FileManager来演示

from contextlib import contextmanager

@contextmanager
def file_manager(name,mode):
    try:
        f = open(name,mode)
        yield f
    finally:
        f.close()

with file_manager('test.txt','w') as f:
    f.write('hello world')

这段代码中，函数file_manager()是一个生成器，当我们执行with语句的时候，便会打开文件，并返回文件对象f，当with语句执行完毕以后，finally代码段中的关闭操作就会执行。

可以看到，使用基于生成器的上下文管理器的时候，我们不用再定义__enter__()和__exit__()两个函数，但是必须加上装饰器，这一点非常容易漏掉。

讲完这两种上下文管理器以后，我们要强调一点：不论是基于类的还是生成器的上下文管理器，两者在功能上是一样的，只不过有下面两点：

1.基于类的上下文管理器更加灵活，适用于大型的系统开发

2.基于生成器的上下文管理器更加方便、简洁，适用于中小型程序。

但是无论使用哪一种，我们一定要记得在exit函数或finally里写好释放资源的代码，这一点尤为重要。

总结

在这一章的开头我们通过一个简单的例子了解了资源泄漏的易发生的特性和其带来的后果，从而引入了上下文管理器这个概念：

　　上下文管理器通常用在文件的IO操作和数据库的连接关闭等场景中，可以确保用过的资源得到迅速释放，有效提高了程序的安全性。

接着，我们通过自定义上下文管理器的示例，大致了解了上下文管理器工作的原理，并介绍基于类的上下文管理器和基于生成器的上下文管理器：两者功能相同，具体使用哪个要根据场景来选择。

另外，上下文管理器通常和with一起使用，大大提高了程序的简洁度，需要注意的是我们在使用with语句执行上下文操作的时候，一旦有异常抛出，异常的类型、值等拘役信息都会通过参数传递给__exit__()函数，我们可以自行定义相关的操作，而在对异常处理完毕以后，务必加上return True语句来保证程序的执行，否则仍然会抛出异常。