Python上下文管理器

上下文管理器

在使用Python编程中，可以会经常碰到这种情况：有一个特殊的语句块，在执行这个语句块之前需要先执行一些准备动作；当语句块执行完成后，需要继续执行一些收尾动作。

例如：当需要操作文件或数据库的时候，首先需要获取文件句柄或者数据库连接对象，当执行完相应的操作后，需要执行释放文件句柄或者关闭数据库连接的动作。

又如，当多线程程序需要访问临界资源的时候，线程首先需要获取互斥锁，当执行完成并准备退出临界区的时候，需要释放互斥锁。

对于这些情况，Python中提供了上下文管理器（Context Manager）的概念，可以通过上下文管理器来定义/控制代码块执行前的准备动作，以及执行后的收尾动作。

上下文管理协议

那么在Python中怎么实现一个上下文管理器呢？这里，又要提到两个"魔术方法"，__enter__和__exit__，下面就是关于这两个方法的具体介绍。

• __enter__(self) Defines what the context manager should do at the beginning of the block created by the with statement. Note that the return value of __enter__ is bound to the target of the with statement, or the name after the as.
• __exit__(self, exception_type, exception_value, traceback) Defines what the context manager should do after its block has been executed (or terminates). It can be used to handle exceptions, perform cleanup, or do something always done immediately after the action in the block. If the block executes successfully, exception_type, exception_value, and traceback will be None. Otherwise, you can choose to handle the exception or let the user handle it; if you want to handle it, make sure __exit__ returns True after all is said and done. If you don't want the exception to be handled by the context manager, just let it happen.

也就是说，当我们需要创建一个上下文管理器类型的时候，就需要实现__enter__和__exit__方法，这对方法就称为上下文管理协议（Context Manager Protocol），定义了一种运行时上下文环境。

with语句

在Python中，可以通过with语句来方便的使用上下文管理器，with语句可以在代码块运行前进入一个运行时上下文（执行__enter__方法），并在代码块结束后退出该上下文（执行__exit__方法）。

with语句的语法如下：

with context_expr [as var]:
    with_suite

• context_expr是支持上下文管理协议的对象，也就是上下文管理器对象，负责维护上下文环境
• as var是一个可选部分，通过变量方式保存上下文管理器对象
• with_suite就是需要放在上下文环境中执行的语句块

在Python的内置类型中，很多类型都是支持上下文管理协议的，例如file，thread.LockType，threading.Lock等等。这里我们就以file类型为例，看看with语句的使用。

with语句简化文件操作

当需要写一个文件的时候，一般都会通过下面的方式。代码中使用了try-finally语句块，即使出现异常，也能保证关闭文件句柄。

logger = open("log.txt", "w")
try:
    logger.write('Hello ')
    logger.write('World')
finally:
    logger.close()

print logger.closed    

其实，Python的内置file类型是支持上下文管理协议的，可以直接通过内建函数dir()来查看file支持的方法和属性：

>>> print dir(file)
['__class__', '__delattr__', '__doc__', '__enter__', '__exit__', '__format__', '
__getattribute__', '__hash__', '__init__', '__iter__', '__new__', '__reduce__',
'__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclass
hook__', 'close', 'closed', 'encoding', 'errors', 'fileno', 'flush', 'isatty', '
mode', 'name', 'newlines', 'next', 'read', 'readinto', 'readline', 'readlines',
'seek', 'softspace', 'tell', 'truncate', 'write', 'writelines', 'xreadlines']
>>>

所以，可以通过with语句来简化上面的代码，代码的效果是一样的，但是使用with语句的代码更加的简洁：

with open("log.txt", "w") as logger:
    logger.write('Hello ')
    logger.write('World')
    
print logger.closed    

自定义上下文管理器

对于自定义的类型，可以通过实现__enter__和__exit__方法来实现上下文管理器。

看下面的代码，代码中定义了一个MyTimer类型，这个上下文管理器可以实现代码块的计时功能：

import time

class MyTimer(object):
    def __init__(self, verbose = False):
        self.verbose = verbose
        
    def __enter__(self):
        self.start = time.time()
        return self
        
    def __exit__(self, *unused):
        self.end = time.time()
        self.secs = self.end - self.start
        self.msecs = self.secs * 1000
        if self.verbose:
            print "elapsed time: %f ms" %self.msecs

下面结合with语句使用这个上下文管理器：

def fib(n):
    if n in [1, 2]:
        return 1
    else:
        return fib(n-1) + fib(n-2)
        
with MyTimer(True):
    print fib(30)

代码输出结果为：

异常处理和__exit__

在使用上下文管理器中，如果代码块 （with_suite）产生了异常，__exit__方法将被调用，而__exit__方法又会有不同的异常处理方式。

当__exit__方法退出当前运行时上下文时，会并返回一个布尔值，该布尔值表明了"如果代码块 （with_suite）执行中产生了异常，该异常是否须要被忽略"。

1. __exit__返回False，重新抛出(re-raised)异常到上层

修改前面的例子，在MyTimer类型中加入了一个参数"ignoreException"来表示上下文管理器是否会忽略代码块 （with_suite）中产生的异常。

import time

class MyTimer(object):
    def __init__(self, verbose = False, ignoreException = False):
        self.verbose = verbose
        self.ignoreException = ignoreException
        
    def __enter__(self):
        self.start = time.time()
        return self
        
    def __exit__(self, *unused):
        self.end = time.time()
        self.secs = self.end - self.start
        self.msecs = self.secs * 1000
        if self.verbose:
            print "elapsed time: %f ms" %self.msecs
        return self.ignoreException
        
try:        
    with MyTimer(True, False):
        raise Exception("Ex4Test")
except Exception, e:
    print "Exception (%s) was caught" %e
else:
    print "No Exception happened"

运行这段代码，会得到以下结果，由于__exit__方法返回False，所以代码块 （with_suite）中的异常会被继续抛到上层代码。

2. __exit__返回Ture，代码块 （with_suite）中的异常被忽略

将代码改为__exit__返回为True的情况：

try:        
    with MyTimer(True, True):
        raise Exception("Ex4Test")
except Exception, e:
    print "Exception (%s) was caught" %e
else:
    print "No Exception happened"

运行结果就变成下面的情况，代码块 （with_suite）中的异常被忽略了，代码继续运行：

一定要小心使用__exit__返回Ture的情况，除非很清楚为什么这么做。

3. 通过__exit__函数完整的签名获取更多异常信息

对于__exit__函数，它的完整签名如下，也就是说通过这个函数可以获得更多异常相关的信息。

• __exit__(self, exception_type, exception_value, traceback)

继续修改上面例子中的__exit__函数如下：

def __exit__(self, exception_type, exception_value, traceback):
    self.end = time.time()
    self.secs = self.end - self.start
    self.msecs = self.secs * 1000
    if self.verbose:
        print "elapsed time: %f ms" %self.msecs
        
    print "exception_type: ", exception_type
    print "exception_value: ", exception_value
    print "traceback: ", traceback
    
    return self.ignoreException

这次运行结果中，就显示出了更多异常相关的信息了：

总结

本文介绍了Python中的上下文管理器，以及如何结合with语句来使用上下文管理器。

总结一下with 语句的执行流程：

• 执行context_expr 以获取上下文管理器对象
• 调用上下文管理器的 __enter__() 方法
  • 如果有 as var 从句，则将 __enter__() 方法的返回值赋给 var
• 执行代码块 with_suite
• 调用上下文管理器的 __exit__() 方法，如果 with_suite 产生异常，那么该异常的 type、value 和 traceback 会作为参数传给 __exit__()，否则传三个 None
  • 如果 with_suite 产生异常，并且 __exit__() 的返回值等于 False，那么这个异常将被重新抛出到上层
  • 如果 with_suite 产生异常，兵器 __exit__() 的返回值等于 True，那么这个异常就被忽略，继续执行后面的代码

在很多情况下，with语句可以简化代码，并增加代码的健壮性。