要自定义一个上下文管理器需要实现__enter__()和__exit__()方法。
如下示例中,创建一具socket网络连接
from socket import socket, AF_INET, SOCK_STREAMclass 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这个类在__enter__方法中建立了一个socket网络连接,在__exit__方法中关闭socket网络连接。在使用时可以通过with语句来自动完成
from functools import partialconn = LazyConnection(('www.python.org', 80))# Connection closedwith 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'')) # conn.__exit__() executes: connection closed说明上下文管理器的原理是将代码会放到 with 语句块中执行。 当出现 with 语句的时候,对象的 __enter__() 方法被触发, 它返回的值(如果有的话)会被赋值给 as 声明的变量。然后,with 语句块里面的代码开始执行。 最后,__exit__() 方法被触发进行清理工作。
不管 with 代码块中发生什么,上面的控制流都会执行完,就算代码块中发生了异常也是一样的。 事实上,__exit__() 方法的三个参数包含了异常类型、异常值和追溯信息(如果有的话)。 __exit__() 方法能自己决定怎样利用这个异常信息,或者忽略它并返回一个None值。 如果 __exit__() 返回 True 那么异常会被清空,就好像什么都没发生一样, with 语句后面的程序继续在正常执行。
还有一个细节问题就是 LazyConnection 是否允许这么多个 with 语句来嵌套使用连接器。 很显然,上面的定义中一次只能允许一个socket连接,如果正在使用一个socket的时候又重复使用 with 语句, 就会产生一个异常了。不过你可以像下面这样修改下上面的实现来解决这个问题:
from socket import socket, AF_INET, SOCK_STREAMclass 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): 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 usefrom functools import partialconn = LazyConnection(('www.python.org', 80))with conn as s1: pass with conn as s2: pass # s1 and s2 are independent sockets上例中LazyConnection 类可以被看做是一个连接工厂。在内部,一个列表被用来构造一个栈。 每次 __enter__() 方法执行的时候,它复制创建一个新的连接并将其加入到栈里面。 __exit__() 方法简单的从栈中弹出最后一个连接并关闭它。
上下文管理器使用在需要管理一些资源比如文件、网络连接和锁的编程环境中,这些资源的一个主要特征是它们必须被手动关闭或释放来确保程序的正确运行。 例如,如果你请求了一个锁,那么你必须确保之后释放了它,否则就可能产生死锁。 通过实现 __enter__() 和 __exit__() 方法并使用 with 语句可以很容易避免这些问题, 因为 __exit__() 方法可以让你无需担心这些了。