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("链接关闭") |