Repository: urllib3/urllib3

Issue:
Add Context Managers
The classes like `PoolManager` and `ConnectionPool` and such should act as context managers to make it easier to clean up the resources used by them.



Discussion:
Makes sense, +1.


Resolving patch:
diff --git a/urllib3/connectionpool.py b/urllib3/connectionpool.py
--- a/urllib3/connectionpool.py
+++ b/urllib3/connectionpool.py
@@ -72,6 +72,21 @@ def __str__(self):
         return '%s(host=%r, port=%r)' % (type(self).__name__,
                                          self.host, self.port)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        # Return False to re-raise any potential exceptions
+        return False
+
+    def close():
+        """
+        Close all pooled connections and disable the pool.
+        """
+        pass
+
+
 # This is taken from http://hg.python.org/cpython/file/7aaba721ebc0/Lib/socket.py#l252
 _blocking_errnos = set([errno.EAGAIN, errno.EWOULDBLOCK])
 
diff --git a/urllib3/poolmanager.py b/urllib3/poolmanager.py
--- a/urllib3/poolmanager.py
+++ b/urllib3/poolmanager.py
@@ -64,6 +64,14 @@ def __init__(self, num_pools=10, headers=None, **connection_pool_kw):
         self.pools = RecentlyUsedContainer(num_pools,
                                            dispose_func=lambda p: p.close())
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.clear()
+        # Return False to re-raise any potential exceptions
+        return False
+
     def _new_pool(self, scheme, host, port):
         """
         Create a new :class:`ConnectionPool` based on host, port and scheme.


Test patch:
diff --git a/test/test_connectionpool.py b/test/test_connectionpool.py
--- a/test/test_connectionpool.py
+++ b/test/test_connectionpool.py
@@ -205,6 +205,26 @@ def test_pool_timeouts(self):
     def test_no_host(self):
         self.assertRaises(LocationValueError, HTTPConnectionPool, None)
 
+    def test_contextmanager(self):
+        with connection_from_url('http://google.com:80') as pool:
+            # Populate with some connections
+            conn1 = pool._get_conn()
+            conn2 = pool._get_conn()
+            conn3 = pool._get_conn()
+            pool._put_conn(conn1)
+            pool._put_conn(conn2)
+
+            old_pool_queue = pool.pool
+
+        self.assertEqual(pool.pool, None)
+
+        self.assertRaises(ClosedPoolError, pool._get_conn)
+
+        pool._put_conn(conn3)
+
+        self.assertRaises(ClosedPoolError, pool._get_conn)
+
+        self.assertRaises(Empty, old_pool_queue.get, block=False)
 
 
 if __name__ == '__main__':
diff --git a/test/test_poolmanager.py b/test/test_poolmanager.py
--- a/test/test_poolmanager.py
+++ b/test/test_poolmanager.py
@@ -71,6 +71,22 @@ def test_nohost(self):
         self.assertRaises(LocationValueError, p.connection_from_url, 'http://@')
         self.assertRaises(LocationValueError, p.connection_from_url, None)
 
+    def test_contextmanager(self):
+        with PoolManager(1) as p:
+            conn_pool = p.connection_from_url('http://google.com')
+            self.assertEqual(len(p.pools), 1)
+            conn = conn_pool._get_conn()
+
+        self.assertEqual(len(p.pools), 0)
+
+        self.assertRaises(ClosedPoolError, conn_pool._get_conn)
+
+        conn_pool._put_conn(conn)
+
+        self.assertRaises(ClosedPoolError, conn_pool._get_conn)
+
+        self.assertEqual(len(p.pools), 0)
+
 
 if __name__ == '__main__':
     unittest.main()