From e60860e65d06d2b4452b7ea94902d79eed11d78c Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Fri, 3 Jun 2016 12:06:36 +1200
Subject: Make tcp.Client.connect return a context manager that closes the
 connection

---
 netlib/tcp.py | 8 ++++++++
 1 file changed, 8 insertions(+)

(limited to 'netlib')

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)
-- 
cgit v1.2.3