32. 让对象支持上下文管理

例如，实现了一个telnet客户端的类TelnetClient，调用实例的 connect() 、 login() 、 interact() 方法启动客户端与服务器交互，交互完毕后调用 cleanup() 方法关闭已连接的socket，以及将操作历史记录写入文件并关闭。

要求：让TelnetClient的实例支持上下文管理协议，从而替代手动调用 connect() 、 cleanup() 方法。

解决方案：

实现上下文管理协议，即实现类的 __enter__() 、 __exit__() 方法，它们分别在with开始和结束时别调用。

对于 with ... as ... 语句：

with所求值的对象必须有一个 __enter__() 方法和一个 __exit__() 方法。

紧跟with后面的语句被求值后，返回对象的 __enter__() 方法被调用，这个方法的返回值将被赋值给as后面的变量。当with后面的代码块全部被执行完之后，将调用前面返回对象的 __exit__() 方法。

方案示例：

# yum install -y telnet-server# systemctl start telnet.socket

from sys import stdin, stdoutimport getpassimport telnetlibfrom collections import dequeclass TelnetClient:
    def __init__(self, host, port=23):
        self.host = host
        self.port = port    def connect(self):
        self.tn = telnetlib.Telnet(self.host, self.port)
        self.history = deque([])

    def cleanup(self):
        self.tn.close()
        self.tn = None

        with open('history.txt', 'a') as f:
            f.writelines(self.history)

    def login(self):
        # user
        self.tn.read_until(b"login: ")
        user = input("Enter your remote account: ")
        self.tn.write(user.encode('utf8') + b"\n")

        # password
        self.tn.read_until(b"Password: ")
        password = getpass.getpass()
        self.tn.write(password.encode('utf8') + b"\n")
        out = self.tn.read_until(b'$ ')
        stdout.write(out.decode('utf8'))

    def interact(self):
        while True:
            cmd = stdin.readline()
            if not cmd:
                break

            self.history.append(cmd)
            self.tn.write(cmd.encode('utf8'))
            out = self.tn.read_until(b'$ ').decode('utf8')

            stdout.write(out[len(cmd)+1:])
            stdout.flush()client = TelnetClient('192.168.30.128')client.connect()client.login()client.interact()client.cleanup()

上面是手工调用 connect() 及 cleanup() 方法。使用 with ... as ... 语句上下文管理：

from sys import stdin, stdoutimport getpassimport telnetlibfrom collections import dequeclass TelnetClient:
    def __init__(self, host, port=23):
        self.host = host
        self.port = port 

    def __enter__(self):
        self.tn = telnetlib.Telnet(self.host, self.port)
        self.history = deque([])
        return self    def __exit__(self, exc_type, exc_value, exc_tb):
        print('IN __exit__', exc_type, exc_value, exc_tb)

        self.tn.close()
        self.tn = None

        with open('history.txt', 'a') as f:
            f.writelines(self.history)

        return True             #将错误压制在方法内部，不再抛给上层

    def login(self):
        # user
        self.tn.read_until(b"login: ")
        user = input("Enter your remote account: ")
        self.tn.write(user.encode('utf8') + b"\n")

        # password
        self.tn.read_until(b"Password: ")
        password = getpass.getpass()
        self.tn.write(password.encode('utf8') + b"\n")
        out = self.tn.read_until(b'$ ')
        stdout.write(out.decode('utf8'))

    def interact(self):
        while True:
            cmd = stdin.readline()
            if not cmd:
                break

            self.history.append(cmd)
            self.tn.write(cmd.encode('utf8'))
            out = self.tn.read_until(b'$ ').decode('utf8')

            stdout.write(out[len(cmd)+1:])
            stdout.flush()with TelnetClient('192.168.30.128') as client:
    raise Exception('TEST')             #生成错误信息
    client.login()
    client.interact()print('END')

with真正强大之处是它可以处理异常。可能你已经注意到类的 __exit__() 方法有三个参数：exc_type, exc_value, exc_tb。在with后面的代码块抛出任何异常时， __exit__() 方法被执行。

查看更多关于32. 让对象支持上下文管理的详细内容...