Python类与对象学习心得-1：让对象支持上下文管理协议

为了让一个对象兼容with 语句，你需要实现它的__enter__() 和__exit__() 方法。
例如，考虑如下的一个类，它能为我们创建一个网络连接：

from socket import socket, AF_INET, SOCK_STREAM

class LazyConnection:
    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.sock = None

    def __enter__(self):
        if self.sock is not None:
            raise RuntimeError('Already connected')
        self.sock = socket(self.family, self.type)
        self.sock.connect(self.address)
        return self.sock

    def __exit__(self, exc_ty, exc_val, tb):
        self.sock.close()
        self.sock = None

这个类的关键特点在于它表示了一个网络连接，但是初始化的时候并不会做任何事情(比如它并没有建立一个连接)，只是记录了socket 的地址和类型信息。连接的建立和关闭是使用 with 语句自动完成的，例如：

from functools import partial

conn = LazyConnection(('www.python.org', 80))
# Connection is closed now
with conn as s:
    # conn.__enter__() executes: connection open
    s.send(b'GET /index.html HTTP/1.0\r\n')
    s.send(b'Host: www.python.org\r\n')
    s.send(b'\r\n')
    resp = b''.join(iter(partial(s.recv, 8192), b'')) # 注意这里是 iter 和 partial 的经典用法
    # conn.__exit__() executes: connection closed

编写上下文管理器的主要原理是你的代码会放到 with 语句块中执行。当出现 with 语句的时候，对象的 __enter__() 方法被触发，它返回的值(如果有的话) 会被赋值给 as 声明的变量。然后，with 语句块里面的代码开始执行。最后，__exit__() 方法被触
发进行清理工作。

不管 with 代码块中发生了什么，上面的控制流都会执行完（即 对象的 __exit__() 方法一定会被触发执行），就算代码块中发生了异常也是一样的。事实上，__exit__() 方法的三个参数包含了异常类型、异常值和追溯信息(如果有的话)。__exit__() 方法能自己决定怎样利用这个异常信息，或者忽略它并返回一个 None 值（如上面的 LazyConnection 类所示，将可能的异常交给上一级上下文处理）。如果 __exit__() 返回 True ，那么所有异常会被清空，就好像什么都没发生一样，with 语句后面的程序继续正常执行。这一点类似于 Python 的 try/finally 机制。

还有一个细节问题就是 LazyConnection 类是否允许多个 with 语句来嵌套使用连接。很显然，上面的定义中一次只能允许一个 socket 连接，如果正在使用一个 socket 的时候又重复使用 with 语句，就会产生一个异常了。不过你可以像下面这样修改上面的实现来解决这个问题：

from socket import socket, AF_INET, SOCK_STREAM

class LazyConnection:
    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.connections = []

    def __enter__(self):
        """创建一个新的socket对象，执行socket连接操作，并加入到self.connections列表中"""
        sock = socket(self.family, self.type)
        sock.connect(self.address)
        self.connections.append(sock)
        return sock

    def __exit__(self, exc_ty, exc_val, tb):
        self.connections.pop().close()


# Example use
conn = LazyConnection(('www.python.org', 80))
with conn as s1:
    pass
    with conn as s2:
        pass
        # s1 and s2 are independent sockets

在第二个版本中，LazyConnection 类可以被看做是某个连接工厂。在内部，一个列表被用来构造一个栈。每次 __enter__() 方法执行的时候，它创建一个新的连接并将其加入到栈里面。__exit__() 方法简单的从栈中弹出最后一个连接并关闭它。这里稍微有点难理解，不过它能允许嵌套使用 with 语句创建多个连接，就如上面演示的那样。

在需要管理一些资源比如文件、网络连接和锁的编程环境中，使用上下文管理器是很普遍的。这些资源的一个主要特征是它们必须被手动的关闭或释放来确保程序的正确运行。例如，如果你请求了一个锁，那么你必须确保之后释放了它，否则就可能产生死锁。通过实现 __enter__() 和 __exit__() 方法并使用 with 语句可以很容易的避免这些问题，因为 __exit__() 方法可以让你无需担心这些了。

上述的为对象单独实现 __enter__() 和 __exit__() 方法还是有些繁琐。实现一个新的上下文管理器的最简单方法是使用 contexlib 模块中的 @contextmanager 装饰器。下面是一个实现了代码块计时功能的上下文管理器例子：

import time
from contextlib import contextmanager

@contextmanager
def timethis(label):
    start = time.time()
    try:
        yield
    finally:
        end = time.time()
        print('{}: {}'.format(label, end - start))


# Example use
with timethis('counting'):
    n = 10000000
    while n > 0:
        n -= 1

在函数 timethis() 中，yield 之前的代码会在上下文管理器中作为 __enter__() 方法执行，所有在 yield 之后的代码会作为__exit__() 方法执行。如果 with 代码块执行时出现了任何异常，异常会在 yield 语句那里抛出。这就是为什么 yield 语句要置于 try 代码块中，同时 yield 之后的代码要置于 finally 代码块中以保证会被执行到（无论是否是否抛出异常）。

如果不使用 contextmanager 装饰器，通常情况下，如果要写一个上下文管理器，你需要定义一个类，里面包含一个 __enter__() 和  __exit__() 方法，如下所示：

import time

class timethis:
    def __init__(self, label):
        self.label = label

    def __enter__(self):
        self.start = time.time()

    def __exit__(self, exc_ty, exc_val, exc_tb):
        end = time.time()
        print('{}: {}'.format(self.label, end - self.start))

尽管这个也不难写，但是和写一个简单的使用 @contextmanager 装饰器的函数相比较还是稍显乏味。

下面是一个更加高级一点的上下文管理器，实现了列表对象上的某种事务(transaction)：

@contextmanager
def list_transaction(orig_list):
    # 进入 with 语句块时为 orig_list 做一个备份列表 working，此后所有修改在 working 列表上进行
    working = list(orig_list) 
    yield working
    # 退出 with 语句块时，如果之前的操作没有抛出异常，则将 working 列表的内容重新覆盖 orig_list
    orig_list[:] = working

这段代码的作用是任何对列表的修改只有当所有代码运行完成并且不出现异常的情况下才会生效。下面我们来演示一下：

>>> items = [1, 2, 3]
>>> with list_transaction(items) as working:
...     working.append(4)
...     working.append(5)
...
>>> items
[1, 2, 3, 4, 5]
>>> with list_transaction(items) as working:
...     working.append(6)
...     working.append(7)
...     raise RuntimeError('oops') # 执行过程中抛出异常，则原有的 items 列表保持原样
...
Traceback (most recent call last):
    File "<stdin>", line 4, in <module>
RuntimeError: oops
>>> items
[1, 2, 3, 4, 5]
>>>

@contextmanager 应该仅仅用来写自包含的上下文管理函数。如果你有一些对象(比如一个文件、网络连接或锁)，需要支持 with 语句，那么你就需要单独实现 __enter__() 方法和 __exit__() 方法。