Merge pull request #953 from heyman/fix_session_safe_mode

Fixes so that safe_mode works for Sessions

AUTHORS.rst

@@ -116,3 +116,4 @@ Patches and Suggestions
 - André Graf (dergraf)
 - Stephen Zhuang (everbird)
 - Martijn Pieters
+- Jonatan Heyman

requests/api.py

@@ -12,10 +12,8 @@ This module implements the Requests API.
 """
 
 from . import sessions
-from .safe_mode import catch_exceptions_if_in_safe_mode
 
 
-@catch_exceptions_if_in_safe_mode
 def request(method, url, **kwargs):
     """Constructs and sends a :class:`Request <Request>`.
     Returns :class:`Response <Response>` object.

requests/safe_mode.py

@@ -18,17 +18,17 @@ import socket
 
 
 def catch_exceptions_if_in_safe_mode(function):
-    """New implementation of safe_mode. We catch all exceptions at the API level
+    """New implementation of safe_mode. We catch all exceptions at the Session level
     and then return a blank Response object with the error field filled. This decorator
-    wraps request() in api.py.
+    wraps Session._send_request() in sessions.py.
     """
 
-    def wrapped(method, url, **kwargs):
+    def wrapped(*args, **kwargs):
         # if save_mode, we catch exceptions and fill error field
         if (kwargs.get('config') and kwargs.get('config').get('safe_mode')) or (kwargs.get('session')
                                             and kwargs.get('session').config.get('safe_mode')):
             try:
-                return function(method, url, **kwargs)
+                return function(*args, **kwargs)
             except (RequestException, ConnectionError, HTTPError,
                     socket.timeout, socket.gaierror) as e:
                 r = Response()
@@ -36,5 +36,5 @@ def catch_exceptions_if_in_safe_mode(function):
                 r.raw = HTTPResponse()  # otherwise, tests fail
                 r.status_code = 0  # with this status_code, content returns None
                 return r
-        return function(method, url, **kwargs)
+        return function(*args, **kwargs)
     return wrapped

requests/sessions.py

@@ -17,6 +17,7 @@ from .models import Request
 from .hooks import dispatch_hook
 from .utils import header_expand, from_key_val_list
 from .packages.urllib3.poolmanager import PoolManager
+from .safe_mode import catch_exceptions_if_in_safe_mode
 
 
 def merge_kwargs(local_kwarg, default_kwarg):
@@ -265,7 +266,12 @@ class Session(object):
             return r
 
         # Send the HTTP Request.
-        r.send(prefetch=prefetch)
+        return self._send_request(r, **args)
+
+    @catch_exceptions_if_in_safe_mode
+    def _send_request(self, r, **kwargs):
+        # Send the request.
+        r.send(prefetch=kwargs.get("prefetch"))
 
         # Return the response.
         return r.response

tests/test_requests.py

@@ -930,6 +930,19 @@ class RequestsTestSuite(TestSetup, TestBaseMixin, unittest.TestCase):
         ds2 = pickle.loads(pickle.dumps(requests.session(prefetch=False)))
         self.assertTrue(ds1.prefetch)
         self.assertFalse(ds2.prefetch)
+    
+    def test_session_connection_error_with_safe_mode(self):
+        config = {"safe_mode":True}
+
+        s = requests.session()
+        r = s.get("http://localhost:1/nope", timeout=0.1, config=config)
+        self.assertFalse(r.ok)
+        self.assertTrue(r.content is None)
+
+        s2 = requests.session(config=config)
+        r2 = s2.get("http://localhost:1/nope", timeout=0.1)
+        self.assertFalse(r2.ok)
+        self.assertTrue(r2.content is None)
 
     def test_connection_error(self):
         try: