python上下文管理器(context manager)

首先，什么是上下文管理器？上下文管理器就是实现了上下文管理协议的对象。主要用于保存和恢复各种全局状态，关闭文件等，上下文管理器本身就是一种装饰器。
下面我们就通过一下几个部分，来介绍下上下文管理器。

上下文管理协议(context management protocol)

上下文管理协议包括两个方法：

• contextmanager.__enter__() 从该方法进入运行时上下文，并返回当前对象或者与运行时上下文相关的其他对象。如果with语句有as关键词存在，返回值会绑定在as后的变量上。

• contextmanager.__exit__(exc_type, exc_val, exc_tb) 退出运行时上下文，并返回一个布尔值标示是否有需要处理的异常。如果在执行with语句体时发生异常，那退出时参数会包括异常类型、异常值、异常追踪信息，否则，3个参数都是None。

我们看一个常见的文件打开操作

>>> with open("/etc/hosts", "r") as file:
...     dir(file)
... 
['__class__', '__delattr__', '__doc__', '__enter__', '__exit__', '__format__', '__getattribute__', '__hash__', '__init__', '__iter__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__', 'close', 'closed', 'encoding', 'errors', 'fileno', 'flush', 'isatty', 'mode', 'name', 'newlines', 'next', 'read', 'readinto', 'readline', 'readlines', 'seek', 'softspace', 'tell', 'truncate', 'write', 'writelines', 'xreadlines']

open函数会返回一个文件类型变量，这个文件类实现了上下文管理协议。

with语句

说上下文管理器就不得不说with语句，为什么呢？
因为with语句就是为支持上下文管理器而存在的，使用上下文管理协议的方法包裹一个代码块（with语句体）的执行，并为try...except...finally提供了一个方便使用的封装。
with语句的语法如下：

with EXPR as VAR:
    BLOCK

with和as是关键词，EXPR就是上下文表达式，是任意表达式（一个表达式，不是表达式列表），VAR是赋值的目标变量。"as VAR"是可选的。
上述语句的底层实现可以这样描述：

mgr = (EXPR)
exit = type(mgr).__exit__  # 并没有调用
value = type(mgr).__enter__(mgr)
exc = True
try:
    try:
        VAR = value  # 如果有"as VAR"
        BLOCK
    except:
        # 这里会处理异常
        exc = False
        if not exit(mgr, *sys.exc_info()):
            raise
        # 如果__exit__返回值是false，异常将被传播；如果返回值是真，异常将被终止
finally:
    if exc:
        exit(mgr, None, None, None)

这样with语句的执行过程就很清楚了。

1. 执行上下文表达式，获取上下文管理器
2. 加载上下文管理器的__exit__()方法以备后期调用
3. 调用上下文管理器的__enter__()方法
4. 如果with语句有指定目标变量，将从__enter__()方法获取的相关对象赋值给目标变量
5. 执行with语句体
6. 调用上下文管理器的__exit__()方法，如果是with语句体造成的异常退出，那异常类型、异常值、异常追踪信息将被传给__exit__()，否则，3个参数都是None。

也可以将多个表达式组织在一起。

with A() as a, B() as b:
    BLOCK

它等价于

with A() as a:
    with B() as b:
        BLOCK

注：多上下文表达式是从python 2.7开始支持的

自定义上下文管理器

class ContextManager(object):
    def __init__(self):
        print '__init__()'

    def __enter__(self):
        print '__enter__()'
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print "__exit__()"

with ContextManager():
    print "OK, we can do something here~~"
    
#输出
__init__()
__enter__()
OK, we can do something here~~
__exit__()

这个例子里，with语句没有使用as关键词，返回了当前对象。下面我们再看一个返回非当前对象的例子

class InnerContext(object):
    def __init__(self, obj):
        print 'InnerContext.__init__(%s)' % obj

    def do_something(self):
        print 'InnerContext.do_something()'

    def __del__(self):
        print 'InnerContext.__del__()'

class ContextManager(object):
    def __init__(self):
        print 'ContextManager.__init__()'

    def __enter__(self):
        print 'ContextManager.__enter__()'
        return InnerContext(self)

    def __exit__(self, exc_type, exc_val, exc_tb):
        print "ContextManager.__exit__()"

with ContextManager() as obj:
    obj.do_something()
    print "OK, we can do something here~~"
    
#输出
ContextManager.__init__()
ContextManager.__enter__()
InnerContext.__init__(<__main__.ContextManager object at 0x1012f95d0>)
InnerContext.do_something()
OK, we can do something here~~
ContextManager.__exit__()
InnerContext.__del__()

下面再来看下异常处理的例子

class ContextManager(object):
    def __init__(self, flag):
        print 'ContextManager.__init__(%s)' % flag
        self.flag = flag

    def __enter__(self):
        print 'ContextManager.__enter__()'
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print 'ContextManager.__exit__(%s, %s, %s)' % (exc_type, exc_val, exc_tb)
        return self.flag

with ContextManager(True):
    raise RuntimeError('error message handled')

print
with ContextManager(False):
    raise RuntimeError('error message propagated')

#输出
ContextManager.__init__(True)
ContextManager.__enter__()
ContextManager.__exit__(<type 'exceptions.RuntimeError'>, error message handled, <traceback object at 0x10d69dbd8>)

ContextManager.__init__(False)
ContextManager.__enter__()
ContextManager.__exit__(<type 'exceptions.RuntimeError'>, error message propagated, <traceback object at 0x109e0fbd8>)
Traceback (most recent call last):
  File "ContextManager.py", line 19, in <module>
    raise RuntimeError('error message propagated')
RuntimeError: error message propagated

contextlib 模块

对于上下文的管理，python也提供了内建的模块contextlib来实现相同的机制，而且这种通过生成器和装饰器实现的上下文管理器，看起来比with语句和手动实现上下文管理协议更优雅。contextlib提供了3个相关函数，分别用于不同的场景。

contextlib.contextmanager(func)

手动实现上下文管理协议并不难，但我们更常遇到的情况是管理非常简单的上下文，这时就可以使用contextmanager将一个生成器转化成上下文管理器，而不需要再去实现上下文管理协议。

import contextlib

@contextlib.contextmanager
def createContextManager(name):
    print '__enter__ %s' %  name
    yield name
    print '__exit__ %s' % name

with createContextManager('foo') as value:
    print value

#输出
__enter__ foo
foo
__exit__ foo

contextlib.nested(mgr1[, mgr2[, ...]])

有时我们会有同时管理多个上下文的需求,这时候就可以使用nested,该函数可以将多个上下文管理器结合为一个嵌套的上下管理器。

with contextlib.nested(createContextManager('a'),createContextManager('b'),createContextManager('c')) as (a, b, c):
    print a, b, c
    
#输出
__enter__ a
__enter__ b
__enter__ c
a b c
__exit__ c
__exit__ b
__exit__ a

等价于

with createContextManager('a') as a:
    with createContextManager('b') as b:
        with createContextManager('c') as c:
            print a, b, c

#输出
__enter__ a
__enter__ b
__enter__ c
a b c
__exit__ c
__exit__ b
__exit__ a

python 2.7 以后，with已经支持多个上下文直接嵌套操作，所以nested已经不建议使用了。

with createContextManager('a') as a,createContextManager('b') as b,createContextManager('c') as c:
    print a, b, c
    
#输出
__enter__ a
__enter__ b
__enter__ c
a b c
__exit__ c
__exit__ b
__exit__ a

contextlib.closing(thing)

从本文的第一个例子里可以看出，文件类是支持上下文管理协议的，可以直接用with语句，还有一些对象并不支持该协议，但使用的时候又要确保正常退出，这时就可以使用closing创建上下文管理器。它等价于

from contextlib import contextmanager

@contextmanager
def closing(thing):
    try:
        yield thing
    finally:
        thing.close()

但使用closing类似如下:

from contextlib import closing
import urllib

with closing(urllib.urlopen('http://www.python.org')) as page:
    for line in page:
        print line

你不必显示的关闭page, 即使抛出异常也会page也会正常关闭。

注： with语句体(with statement body) 在with语句中，with包裹起来的语句体并没有一个非常正式的名称，只是一般大家都这么叫。

reference

• https://www.python.org/dev/peps/pep-0343/
• https://docs.python.org/2.7/library/stdtypes.html#typecontextmanager
• https://docs.python.org/2.7/reference/compound_stmts.html#with
• https://docs.python.org/2.7/library/contextlib.html?highlight=contextlib#contextlib.contextmanager
• https://pymotw.com/2/contextlib/