58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
|
# 有时候我们想让类支持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
|
|||
|
self.sock = socket(self.family, self.type)
|
|||
|
self.sock.connect(self.address)
|
|||
|
print("链接被拉起")
|
|||
|
return self.sock
|
|||
|
|
|||
|
def __exit__(self, exc_ty, exc_va, tb):
|
|||
|
self.sock.close()
|
|||
|
self.sock = None
|
|||
|
print("链接关闭")
|
|||
|
|
|||
|
from functools import partial
|
|||
|
conn = LazyConnection(('www.python.org', 80))
|
|||
|
with conn as s:
|
|||
|
# 链接被拉起
|
|||
|
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, 1024), b''))
|
|||
|
# 链接被关闭
|
|||
|
|
|||
|
# 但是这个类只能使用单层的with,如果有多个with嵌套,它就会失效
|
|||
|
# 可以对原来的类进行一次修改,做一个工厂类,每次从一个栈中存取一个链接
|
|||
|
class LazyConnectionV2():
|
|||
|
def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
|
|||
|
self.address = address
|
|||
|
self.family = family
|
|||
|
self.type = type
|
|||
|
# 单个的sock被替换成空列表
|
|||
|
self.connections = []
|
|||
|
|
|||
|
def __enter__(self):
|
|||
|
# 建立一个socket链接
|
|||
|
sock = socket(self.family, self.type)
|
|||
|
sock.connect(self.address)
|
|||
|
# 将链接压入栈中
|
|||
|
self.connections.append(sock)
|
|||
|
print("链接被拉起")
|
|||
|
return sock
|
|||
|
|
|||
|
def __exit__(self, exc_ty, exc_va, tb):
|
|||
|
# 将最近被压入栈的链接弹出并关闭
|
|||
|
self.connections.pop().close()
|
|||
|
print("链接关闭")
|