#25485: Add context manager support to Telnet class.
Patch by Stéphane Wirtel.
diff --git a/Lib/telnetlib.py b/Lib/telnetlib.py
index 72dabc7..b0863b1 100644
--- a/Lib/telnetlib.py
+++ b/Lib/telnetlib.py
@@ -637,6 +637,12 @@
raise EOFError
return (-1, None, text)
+ def __enter__(self):
+ return self
+
+ def __exit__(self, type, value, traceback):
+ self.close()
+
def test():
"""Test program for telnetlib.
@@ -660,11 +666,10 @@
port = int(portstr)
except ValueError:
port = socket.getservbyname(portstr, 'tcp')
- tn = Telnet()
- tn.set_debuglevel(debuglevel)
- tn.open(host, port, timeout=0.5)
- tn.interact()
- tn.close()
+ with Telnet() as tn:
+ tn.set_debuglevel(debuglevel)
+ tn.open(host, port, timeout=0.5)
+ tn.interact()
if __name__ == '__main__':
test()
diff --git a/Lib/test/test_telnetlib.py b/Lib/test/test_telnetlib.py
index 524bba3..23029e0 100644
--- a/Lib/test/test_telnetlib.py
+++ b/Lib/test/test_telnetlib.py
@@ -42,6 +42,11 @@
telnet = telnetlib.Telnet(HOST, self.port)
telnet.sock.close()
+ def testContextManager(self):
+ with telnetlib.Telnet(HOST, self.port) as tn:
+ self.assertIsNotNone(tn.get_socket())
+ self.assertIsNone(tn.get_socket())
+
def testTimeoutDefault(self):
self.assertTrue(socket.getdefaulttimeout() is None)
socket.setdefaulttimeout(30)