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, stdout
import getpass
import telnetlib
from collections import deque

class 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, stdout
import getpass
import telnetlib
from collections import deque

class 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__()方法被執行。

---