Make tcp.Client.connect return a context manager that closes the connection

1 files changed, 8 insertions, 0 deletions

diff --git a/netlib/tcp.py b/netlib/tcp.py
index bb0c93a9..61209d64 100644
--- a/netlib/tcp.py
+++ b/netlib/tcp.py
@@ -6,6 +6,7 @@ import sys
 import threading
 import time
 import traceback
+import contextlib
 
 import binascii
 from six.moves import range
@@ -577,6 +578,12 @@ class _Connection(object):
         return context
 
 
+@contextlib.contextmanager
+def _closer(client):
+    yield
+    client.close()
+
+
 class TCPClient(_Connection):
 
     def __init__(self, address, source_address=None):
@@ -708,6 +715,7 @@ class TCPClient(_Connection):
         self.connection = connection
         self.ip_address = Address(connection.getpeername())
         self._makefile()
+        return _closer(self)
 
     def settimeout(self, n):
         self.connection.settimeout(n)