akaIDIOT · March 25, 2016 14:55 · Mar 25, 2016 · Mar 25, 2016 · Mar 25, 2016
diff --git a/drain.py b/drain.py
@@ -28,8 +28,6 @@ def __enter__(self):
         return self.response
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if not self.response._content_consumed:
-            # no need to do anything if the content was already consumed
-            # content is a property that consumes all of response's content on access
-            self.response.content
-        return False
+        # accessing response.content will make sure it is drained
+        self.response.content
+        return False
diff --git a/drain.py b/drain.py
@@ -25,7 +25,7 @@ def __init__(self, response):
         self.response = response
 
     def __enter__(self):
-        return self
+        return self.response
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if not self.response._content_consumed:

diff --git a/drain.py b/drain.py
@@ -0,0 +1,35 @@
+class drain(object):
+    """
+    Context manager to make sure a `Response` object is drained of content.
+    Can be used to return a connection used to retrieve a response back to the
+    connection pool it came from in case the actual response content of the
+    response is not needed, e.g.:
+
+    .. code-block:: python
+
+        response = session.get('https://example.com/')
+
+        with drain(response):
+            # we're only interested in the status code here and ignore the
+            # response content, drain will make sure the content is consumed 
+            # and the underlying connection can be reused
+            return response.status_code == 200
+    """
+
+    def __init__(self, response):
+        """
+        Creates a new draining context manager for *response*.
+
+        :param response: the `Response` to drain on exit
+        """
+        self.response = response
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if not self.response._content_consumed:
+            # no need to do anything if the content was already consumed
+            # content is a property that consumes all of response's content on access
+            self.response.content
+        return False