python核心技术与实战学习笔记（十六）：巧用上下文管理器和with语句精简代码

16.1 上下文管理器简介

在任何一门编程语言中，文件的输入输出、数据库的连接断开等，都是很常见的资源管理操作。但是资源是有限的，我们必须注意要在使用资源后进行资源的释放，否则容易造成资源泄露，使得系统处理缓慢甚至会系统奔溃。

为了解决这种问题，不同的编程语言引入了不同的机制。在python语言中，对应的解决方式便是上下文管理器(contex manager)。上下文管理器可以帮助自动分配并且释放资源，其中最常用的应用便是with语句。这就解释了为什么之前更推荐使用with语句来完成文件的IO处理。

另一个典型的例子是threading.lock类。比如想获得一个锁，执行相应的操作，完成后再释放。那么代码可写为：

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

而对应的with语句，则显得非常简洁：

some_lock = threading.Lock()
with somelock:
    ...

总之，with语句可以使得代码更加简洁，同时有效避免资源泄露的问题。

16.2 上下文管理器的实现

16.2.1 基于类的上下文管理器

接下来通过例子来看看上下文管理器的原理，弄清楚它的实现。这里通过定义一个上下文管理类FileManager，模拟Python的打开、关闭文件操作：

class FileManager:
    def __init__(self, name, mode):
        print('calling __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('calling __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')
    
## 输出
calling __init__ method
calling __enter__ method
ready to write to file
calling __exit__ method

当我们用类来创建上下文管理器时，必须保证这个类包括方法“__enter__()”和方法"__exit__()"：

• __enter__()返回需要被管理的资源
• __exit__()通常存在一些释放、清理资源的操作。如上面代码中的关闭文件等等。

当我们使用with语句执行这个上下文管理器时：

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

执行步骤为：

1. 方法“__init__()”被调用，程序初始化对象FileManager，使得文件名(name)为’test.txt’，文件模式(mode)是‘w’；
2. 方法“__enter__()”被调用，文件以写入模式被打开，并且返回FileManager对象赋予f
3. 字符串“hello”被写入文件“test.txt”；
4. 方法"__exit__()"被调用，负责关闭之前打开的文件流

因此程序输出为：

calling __init__ method
calling __enter__ method
ready to write to file
calling __exit__ meth

此外，还可以在方法"__exit__()“中通过设置参数捕获异常，其参数“exc_type, exc_val, exc_tb”
分别表示exception_type、exception_value和trackback。当我们执行含有上下文管理器的with语句时，如若有异常抛出，异常的信息就会包含在这三个变量中，传入方法”__exit__()"。

如下面代码所示，如需处理可能发生的异常，可以在"__exit__()"添加相应的代码：

class Foo:
    def __init__(self):
        print('__init__ called')        

    def __enter__(self):
        print('__enter__ called')
        return self
    
    def __exit__(self, exc_type, exc_value, exc_tb):
        print('__exit__ called')
        if exc_type:
            print(f'exc_type: {exc_type}')
            print(f'exc_value: {exc_value}')
            print(f'exc_traceback: {exc_tb}')
            print('exception handled')
        return True    # 如果确定异常已被处理，需要加上return True
    
with Foo() as obj:
    raise Exception('exception raised').with_traceback(None)  # 手动抛出异常

# 输出
__init__ called
__enter__ called
__exit__ called
exc_type: <class 'Exception'>
exc_value: exception raised
exc_traceback: <traceback object at 0x1046036c8>
exception handled

同样，对于数据库的连接操作，也常用上下文管理器来表示，简化代码如下：

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', '8080') as db_client: 
	...

与前面FileManager的例子相似：

• 方法“__init__()”负责对数据库进行初始化，也就是将主机名’localhost’和接口 '8080’分别赋予变量hostname和port
• 方法“__enter__()”连接数据库，并且返回对象DBConnectionManager
• 方法"__exit__()"负责关闭数据库的连接

有了DBConnectionManager这个类，那么在程序调用数据库时，我们都只需简单地调用with语句即可，并不需要关心数据库的关闭、异常等等，显然大大提高了开发的效率。

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

基于生成器的上下文管理器和基于类的上下文管理器在功能上是一致的，区别在于：

• 基于类的上下文管理器更加灵活，适用于大型的系统开发
• 基于生成器的上下文管理器更加方便简洁，适用于中小型程序

我们可以使用装饰器contexlib.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 block中的关闭文件操作便会执行。之所以说使用基于生成器的上下文管理器会更简洁方便的原因是，我们无需像基于类的上下文管理器那样定义“__enter__()”和“__exit__()”方法，但一定要使用装饰器@contextmanager ！

最后，要注意的是，无论是使用基于类的上下文管理器和基于生成器的上下文管理器，必须要在方法“__exit__()”或者finally block中释放资源。