Python进阶 上下文管理器(Context managers)

上下文管理器允许你在有需要的时候，精确地分配和释放资源。
使用上下文管理器最广泛的案例就是with语句了。
想象下你有两个需要结对执行的相关操作，然后还要在它们中间放置一段代码。
上下文管理器就是专门让你做这种事情的。

with open('some_file', 'w') as opened_file:
    opened_file.write('Hola!')

# 上面那段代码与这一段是等价的
file = open('some_file', 'w')
try:
    file.write('Hola!')
finally:
    file.close()

基于类的实现

一个上下文管理器的类，最起码要定义enter和exit方法。

class File(object):
    def __init__(self, file_name, method):
        self.file_obj = open(file_name, method)
    def __enter__(self):
        return self.file_obj
    def __exit__(self, type, value, traceback):
        self.file_obj.close()

# 通过定义__enter__和__exit__方法，我们可以在with语句里使用它。
with File('demo.txt', 'w') as opened_file:
    opened_file.write('Hola!')

我们来谈谈在底层都发生了什么。

1. with语句先暂存了File类的exit方法
2. 然后它调用File类的enter方法
3. enter方法打开文件并返回给with语句
4. 打开的文件句柄被传递给opened_file参数
5. 我们使用.write()来写文件
6. with语句调用之前暂存的exit方法
7. exit方法关闭了文件

我们还没有谈到exit方法的这三个参数：type, value和traceback。
在第4步和第6步之间，如果发生异常，Python会将异常的type,value和traceback传递给exit方法。
它让exit方法来决定如何关闭文件以及是否需要其他步骤。在我们的案例中，我们并没有注意它们。
那如果我们的文件对象抛出一个异常呢？万一我们尝试访问文件对象的一个不支持的方法。举个例子：

with File('demo.txt', 'w') as opened_file:
    opened_file.undefined_function('Hola!')

我们来列一下，当异常发生时，with语句会采取哪些步骤。

1. 它把异常的type,value和traceback传递给exit方法
2. 它让exit方法来处理异常
3. 如果exit返回的是True，那么这个异常就被优雅地处理了。
   4.如果exit返回的是True以外的任何东西，那么这个异常将被with语句抛出。

在我们的案例中，exit方法返回的是None(如果没有return语句那么方法会返回None)。因此，with语句抛出了那个异常。

Traceback (most recent call last):
  File "", line 2, in 
AttributeError: 'file' object has no attribute 'undefined_function'

我们尝试下在exit方法中处理异常：

class File(object):
    def __init__(self, file_name, method):
        self.file_obj = open(file_name, method)
    def __enter__(self):
        return self.file_obj
    def __exit__(self, type, value, traceback):
        print("Exception has been handled")
        self.file_obj.close()
        return True

with File('demo.txt', 'w') as opened_file:
    opened_file.undefined_function()

# Output: Exception has been handled

我们的exit方法返回了True,因此没有异常会被with语句抛出。

基于生成器的实现

我们还可以用装饰器(decorators)和生成器(generators)来实现上下文管理器。
Python有个contextlib模块专门用于这个目的。我们可以使用一个生成器函数来实现一个上下文管理器，而不是使用一个类。

from contextlib import contextmanager

@contextmanager
def open_file(name):
    try:
        f = open(name, 'w')
        yield f
    except BaseException as e:
        print(e)
        raise e
    finally:
        f.close()

OK啦！这个实现方式看起来更加直观和简单。然而，这个方法需要关于生成器、yield和装饰器的一些知识。它的工作方式和之前的方法大致相同。

让我们小小地剖析下这个方法。

1. Python解释器遇到了yield关键字。因为这个缘故它创建了一个生成器而不是一个普通的函数。
2. 因为这个装饰器，contextmanager会被调用并传入函数名（open_file）作为参数。
3. contextmanager函数返回一个以GeneratorContextManager对象封装过的生成器。
4. 这个GeneratorContextManager被赋值给open_file函数，我们实际上是在调用GeneratorContextManager对象。

使用这个新生成的上下文管理器：

with open_file('some_file') as f:
    f.write('hola!')