Python标准模块--ContextManager

1 模块简介

在数年前，Python 2.5 加入了一个很是特殊的关键字，就是with。with语句容许开发者建立上下文管理器。什么是上下文管理器？上下文管理器就是容许你能够自动地开始和结束一些事情。例如，你可能想要打开一个文件，而后写入一些内容，最后再关闭文件。这或许就是上下文管理器中一个最经典的示例。事实上，当你利用with语句打开一个文件时，Python替你自动建立了一个上下文管理器。

with open("test/test.txt","w") as f_obj:
    f_obj.write("hello")

若是你使用的是Python 2.4，你不得不以一种老的方式来完成这个任务。

f_obj = open("test/test.txt","w")
f_obj.write("hello")
f_obj.close()

上下文管理器背后工做的机制是使用Python的方法：__enter__和__exit__。让咱们尝试着去建立咱们的上下文管理器，以此来了解上下文管理器是如何工做的。

2 模块使用

2.1 建立一个上下文管理器类

与其继续使用Python打开文件这个例子，不如咱们建立一个上下文管理器，这个上下文管理器将会建立一个SQLite数据库链接，当任务处理完毕，将会将其关闭。下面就是一个简单的示例。

import sqlite3

class DataConn:
    def __init__(self,db_name):
        self.db_name = db_name

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_name)
        return self.conn

    def __exit__(self,exc_type,exc_val,exc_tb):
        self.conn.close()
        if exc_val:
            raise

if __name__ == "__main__":
    db = "test/test.db"
    with DataConn(db) as conn:
        cursor = conn.cursor()

在上述代码中，咱们建立了一个类，获取到SQLite数据库文件的路径。__enter__方法将会自动执行，并返回数据库链接对象。如今咱们已经获取到数据库链接对象，而后咱们建立光标，向数据库写入数据或者对数据库进行查询。当咱们退出with语句的时候，它将会调用__exit__方法用于执行和关闭这个链接。

让咱们使用其它的方法来建立上下文管理器。

2.2 利用contextlib建立一个上下文管理器

Python 2.5 不单单添加了with语句，它也添加了contextlib模块。这就容许咱们使用contextlib的contextmanager函数做为装饰器，来建立一个上下文管理器。让咱们尝试着用它来建立一个上下文管理器，用于打开和关闭文件。

from contextlib import contextmanager

@contextmanager
def file_open(path):
    try:
        f_obj = open(path,"w")
        yield f_obj
    except OSError:
        print("We had an error!")
    finally:
        print("Closing file")
        f_obj.close()

if __name__ == "__main__":
    with file_open("test/test.txt") as fobj:
        fobj.write("Testing context managers")

在这里，咱们从contextlib模块中引入contextmanager，而后装饰咱们所定义的file_open函数。这就容许咱们使用Python的with语句来调用file_open函数。在函数中，咱们打开文件，而后经过yield，将其传递出去，最终主调函数可使用它。

一旦with语句结束，控制就会返回给file_open函数，它继续执行yield语句后面的代码。这个最终会执行finally语句--关闭文件。若是咱们在打开文件时遇到了OSError错误，它就会被捕获，最终finally语句依然会关闭文件句柄。

contextlib.closing(thing)

contextlib模块提供了一些很方便的工具。第一个工具就是closing类，一旦代码块运行完毕，它就会将事件关闭。Python官方文档给出了相似于如下的一个示例，

>>> from contextlib import contextmanager
>>> @contextmanager
... def closing(db):
...     try:
...         yield db.conn()
...     finally:
...         db.close()

在这段代码中，咱们建立了一个关闭函数，它被包裹在contextmanager中。这个与closing类相同。区别就是，咱们能够在with语句中使用closing类自己，而非装饰器。让咱们看以下的示例，

>>> from contextlib import closing
>>> from urllib.request import urlopen
>>> with closing(urlopen("http://www.google.com")) as webpage:
...     for line in webpage:
...         pass

在这个示例中，咱们在closing类中打开一个url网页。一旦咱们运行完毕with语句，指向网页的句柄就会关闭。

contextlib.suppress(*exceptions)

另外一个工具就是在Python 3.4中加入的suppress类。这个上下文管理工具背后的理念就是它能够禁止任意数目的异常。假如咱们想忽略FileNotFoundError异常。若是你书写了以下的上下文管理器，那么它不会正常运行。

>>> with open("1.txt") as fobj:
...     for line in fobj:
...         print(line)
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
FileNotFoundError: [Errno 2] No such file or directory: '1.txt'

正如你所看到的，这个上下文管理器没有处理这个异常，若是你想忽略这个错误，你能够按照以下方式来作，

>>> from contextlib import suppress
>>> with suppress(FileNotFoundError):
...     with open("1.txt") as fobj:
...         for line in fobj:
...             print(line)

在这段代码中，咱们引入suppress，而后将咱们要忽略的异常传递给它，在这个例子中，就是FileNotFoundError。若是你想运行这段代码，你将会注意到，文件不存在时，什么事情都没有发生，也没有错误被抛出。请注意，这个上下文管理器是可重用的，2.4章节将会具体解释。

contextlib.redirect_stdout/redirect_stderr

contextlib模块还有一对用于重定向标准输出和标准错误输出的工具，分别在Python 3.4 和3.5 中加入。在这些工具被加入以前，若是你想对标准输出重定向，你须要按照以下方式操做，

import sys
path = "test/test.txt"

with open(path,"w") as fobj:
    sys.stdout = fobj
    help(sum)

利用contextlib模块，你能够按照以下方式操做，

from contextlib import redirect_stdout

path = "test/test.txt"

with open(path,"w") as fobj:
    with redirect_stdout(fobj):
        help(redirect_stdout)

在上面两个例子中，咱们均是将标准输出重定向到一个文件。当咱们调用Python的help函数，不是将信息输出到标准输出上，而是将信息保存到重定向的文件中。你也能够将标准输出重定向到缓存或者从用接口如Tkinter或wxPython中获取的文件控制类型上。

2.3 ExitStack

ExitStack是一个上下文管理器，容许你很容易地与其它上下文管理结合或者清除。这个咋听起来让人有些迷糊，咱们来看一个Python官方文档的例子，或许会让咱们更容易理解它。

>>> from contextlib import ExitStack
>>> filenames = ["1.txt","2.txt"]
>>> with ExitStack as stack:
...     file_objects = [stack.enter_context(open(filename)) for filename in filenames]

这段代码就是在列表中建立一系列的上下文管理器。ExitStack维护一个寄存器的栈。当咱们退出with语句时，文件就会关闭，栈就会按照相反的顺序调用这些上下文管理器。

Python官方文档中关于contextlib有不少示例，你能够学习到以下的技术点：

• 从__enter__方法中捕获异常
• 支持不定数目的上下文管理器
• 替换掉try-finally
• 其它

2.4 可重用的上下文管理器

大部分你所建立的上下文管理器仅仅只能在with语句中使用一次，示例以下：

>>> from contextlib import contextmanager
>>> @contextmanager
... def single():
...     print("Yielding")
...     yield
...     print("Exiting context manager")
...
>>> context = single()
>>> with context:
...     pass
...
Yielding
Exiting context manager
>>> with context:
...     pass
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3.4/contextlib.py", line 61, in __enter__
    raise RuntimeError("generator didn't yield") from None
RuntimeError: generator didn't yield

在这段代码中，咱们建立了一个上下文管理器实例，并尝试着在Python的with语句中运行两次。当第二次运行时，它抛出了RuntimeError。

可是若是咱们想运行上下文管理器两次呢？咱们须要使用可重用的上下文管理器。让咱们使用以前所用过的redirect_stdout这个上下文管理器做为示例，

>>> from contextlib import redirect_stdout
>>> from io import StringIO
>>> stream = StringIO()
>>> write_to_stream = redirect_stdout(stream)
>>> with write_to_stream:
...     print("Write something to the stream")
...     with write_to_stream:
...         print("Write something else to stream")
...
>>> print(stream.getvalue())
Write something to the stream
Write something else to stream

在这段代码中，咱们建立了一个上下文管理器，它们均向StringIO（一种内存中的文件流）写入数据。这段代码正常运行，而没有像以前那样抛出RuntimeError错误，缘由就是redirect_stdout是可重用的，容许咱们能够调用两次。固然，实际的例子将会有更多的函数调用，会更加的复杂。必定要注意，可重用的上下文管理器不必定是线程安全的。若是你须要在线程中使用它，请先仔细阅读Python的文档。

2.5 总结

上下文管理器颇有趣，也很方便。我常常在自动测试中使用它们，例如，打开和关闭对话。如今，你应该可使用Python内置的工具去建立你的上下文管理器。你还能够继续阅读Python关于contextlib的文档，那里有不少本文没有覆盖到的知识。

3 Reference

Python 201