如何让对象支持上下文管理

问题举例

一个telnet客户端的类TelnetClient, 调用实例的connect(),login(),interact方法

启动客户端与服务器交互,交互完毕后需要调用cleanup()方法,关闭已连接的socket,

以及将操作历史记录写入文件并关闭。

能否让TelnetClient的实例支持上下文管理协议,从而代替手工调用connect(),cleanup()方法。

解决思路

实现上下文管理协议,即实现类的__enter__, __exit__方法,

它们分别在with开始和结束时被调用。

代码

说明:这段代码并不能直接运行不了,因为服务器ip连接不上,

这里大家感受下大致实现思路即可即可

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() # client = TelnetClient('192.168.0.105')
# client.connect()
# client.login()
# client.interact()
# client.cleanup() with TelnetClient('192.168.0.105') as client:
raise Exception('TEST')
client.login()
client.interact() print('END')

__exit__中返回True是为了压制异常向上抛

参考资料:python3实用编程技巧进阶

05-08 15:50