Make hooks work with prepared requests
authorsprt <hellosprt@gmail.com>
Sat, 12 Jan 2013 20:46:44 +0000 (21:46 +0100)
committersprt <hellosprt@gmail.com>
Sat, 12 Jan 2013 20:46:44 +0000 (21:46 +0100)
requests/models.py
requests/sessions.py
test_requests.py

index 5202e6f..0b94262 100644 (file)
@@ -11,7 +11,7 @@ import collections
 import logging
 
 from io import BytesIO
-from .hooks import default_hooks
+from .hooks import default_hooks, HOOKS
 from .structures import CaseInsensitiveDict
 from .status_codes import codes
 
@@ -225,6 +225,7 @@ class Request(RequestHooksMixin):
         # Note that prepare_auth must be last to enable authentication schemes
         # such as OAuth to work on a fully prepared request.
         p.prepare_auth(self.auth)
+        p.prepare_hooks(self.hooks)
 
         return p
 
@@ -415,6 +416,21 @@ class PreparedRequest(RequestEncodingMixin, RequestHooksMixin):
             if cookie_header is not None:
                 self.headers['Cookie'] = cookie_header
 
+    def prepare_hooks(self, hooks):
+        """Prepares the given hooks."""
+        for event in HOOKS:
+            if event not in self.hooks:
+                self.hooks[event] = []
+            if event not in hooks:
+                hooks[event] = []
+
+            if not hasattr(self.hooks[event], '__iter__'):
+                self.hooks[event] = [self.hooks[event]]
+            if not hasattr(hooks[event], '__iter__'):
+                hooks[event] = [hooks[event]]
+
+            self.hooks[event].extend(hooks[event])
+
 
 class Response(object):
     """The :class:`Response <Response>` object, which contains a
index d65877c..7dfabfe 100644 (file)
@@ -13,7 +13,7 @@ import os
 from .compat import cookielib
 from .cookies import cookiejar_from_dict
 from .models import Request
-from .hooks import dispatch_hook, default_hooks
+from .hooks import default_hooks
 from .utils import from_key_val_list, default_headers
 from .exceptions import TooManyRedirects, InvalidSchema
 
@@ -294,9 +294,6 @@ class Session(SessionRedirectMixin):
             resp = history.pop()
             resp.history = tuple(history)
 
-        # Response manipulation hook.
-        self.response = dispatch_hook('response', hooks, resp)
-
         return resp
 
     def get(self, url, **kwargs):
index 4a70982..6790fea 100644 (file)
@@ -264,6 +264,19 @@ class RequestsTestCase(unittest.TestCase):
         self.assertEqual(r.status_code, 200)
         self.assertTrue(b"text/py-content-type" in r.request.body)
 
+    def test_prepared_request_hook(self):
+        def hook(resp):
+            resp.hook_working = True
+            return resp
+
+        req = requests.Request('GET', HTTPBIN, hooks={'response': hook})
+        prep = req.prepare()
+
+        s = requests.Session()
+        resp = s.send(prep)
+
+        self.assertTrue(hasattr(resp, 'hook_working'))
+
 
 if __name__ == '__main__':
     unittest.main()