# 有时候我们想让类支持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("链接关闭")