Files
Python_CookBook_repo/8.类与对象/3.让对象支持上下文管理协议.py

58 lines
1.9 KiB
Python
Raw Permalink Normal View History

2025-09-10 16:12:45 +08:00
# 有时候我们想让类支持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("链接关闭")