问题
让对象支持上下文管理协议,兼容with语句。
解决方案
为了让一个对象兼容 with 语句,需要实现 __enter__() 和__exit__() 方法。 例如,考虑如下的一个类,它能为我们创建一个网络连接:
from socket import socket, AF_INET, SOCK_STREAM
class LazyConnection(object):
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 = sock(self.family, self.type)
self.sock.connected(self.address)
return self.sock
def __exit__(self, exc_ty, exc_val, tb):
self.sock.close()
self.sock = None
这个类的关键特点在于它表示了一个网络连接,但是初始化的时候并不会做任何事情。 连接的建立和关闭是使用 with 语句自动完成的。
讨论
编写上下文管理器的主要原理是将代码会放到 with 语句块中执行。 当出现 with 语句的时候,对象的__enter__()方法被触发, 它返回的值(如果有的话),会被赋值给 as 声明的变量。然后,with 语句块里面的代码开始执行。 最后__exit__() 方法被触发进行清理工作。
不管 with 代码块中发生什么,上面的控制流都会执行完,就算代码块中发生了异常。 事实上,__exit__() 方法的第三个参数包含了异常类型、异常值和追溯信息(如果有的话)。 __exit__() 方法能自己决定怎样利用这个异常信息,或者忽略它并返回一个None值。 如果 __exit__() 返回 True ,那么异常会被清空,就好像什么都没发生, with 语句后面的程序继续正常执行。