[Feature] resource aware hooks

falcon/api_helpers.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import re
-from functools import wraps
 
 from falcon import responders, HTTP_METHODS
 import falcon.status_codes as status
+from falcon.hooks import _wrap_with_hooks
 
 STREAM_BLOCK_SIZE = 8 * 1024  # 8 KiB
 
@@ -217,7 +217,8 @@ def create_http_method_map(resource, uri_fields, before, after):
         else:
             # Usually expect a method, but any callable will do
             if callable(responder):
-                responder = _wrap_with_hooks(before, after, responder)
+                responder = _wrap_with_hooks(
+                    before, after, responder, resource)
                 method_map[method] = responder
 
     # Attach a resource for unsupported HTTP methods
@@ -229,68 +230,15 @@ def create_http_method_map(resource, uri_fields, before, after):
         # OPTIONS itself is intentionally excluded from the Allow header
         responder = responders.create_default_options(
             allowed_methods)
-        method_map['OPTIONS'] = _wrap_with_hooks(before, after, responder)
+        method_map['OPTIONS'] = _wrap_with_hooks(
+            before, after, responder, resource)
         allowed_methods.append('OPTIONS')
 
     na_responder = responders.create_method_not_allowed(allowed_methods)
 
     for method in HTTP_METHODS:
         if method not in allowed_methods:
-            method_map[method] = _wrap_with_hooks(before, after, na_responder)
+            method_map[method] = _wrap_with_hooks(
+                before, after, na_responder, resource)
 
     return method_map
-
-
-# -----------------------------------------------------------------------------
-# Helpers
-# -----------------------------------------------------------------------------
-
-
-def _wrap_with_hooks(before, after, responder):
-    if after is not None:
-        for action in after:
-            responder = _wrap_with_after(action, responder)
-
-    if before is not None:
-        # Wrap in reversed order to achieve natural (first...last)
-        # execution order.
-        for action in reversed(before):
-            responder = _wrap_with_before(action, responder)
-
-    return responder
-
-
-def _wrap_with_before(action, responder):
-    """Execute the given action function before a bound responder.
-
-    Args:
-        action: A function with a similar signature to a resource responder
-            method, taking (req, resp, params).
-        responder: The bound responder to wrap.
-
-    """
-
-    @wraps(responder)
-    def do_before(req, resp, **kwargs):
-        action(req, resp, kwargs)
-        responder(req, resp, **kwargs)
-
-    return do_before
-
-
-def _wrap_with_after(action, responder):
-    """Execute the given action function after a bound responder.
-
-    Args:
-        action: A function with a signature similar to a resource responder
-            method, taking (req, resp).
-        responder: The bound responder to wrap.
-
-    """
-
-    @wraps(responder)
-    def do_after(req, resp, **kwargs):
-        responder(req, resp, **kwargs)
-        action(req, resp)
-
-    return do_after

falcon/hooks.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from functools import wraps
+import inspect
+
 import six
 
 from falcon import HTTP_METHODS
@@ -60,10 +62,8 @@ def _before(responder_or_resource):
                         # variable that is shared between iterations of the
                         # for loop, above.
                         def let(responder=responder):
-                            @wraps(responder)
-                            def do_before_all(self, req, resp, **kwargs):
-                                action(req, resp, kwargs)
-                                responder(self, req, resp, **kwargs)
+                            do_before_all = _wrap_with_before(
+                                action, responder, resource, True)
 
                             setattr(resource, responder_name, do_before_all)
 
@@ -73,11 +73,7 @@ def do_before_all(self, req, resp, **kwargs):
 
         else:
             responder = responder_or_resource
-
-            @wraps(responder)
-            def do_before_one(self, req, resp, **kwargs):
-                action(req, resp, kwargs)
-                responder(self, req, resp, **kwargs)
+            do_before_one = _wrap_with_before(action, responder, None, True)
 
             return do_before_one
 
@@ -107,11 +103,10 @@ def _after(responder_or_resource):
                 else:
                     # Usually expect a method, but any callable will do
                     if callable(responder):
+
                         def let(responder=responder):
-                            @wraps(responder)
-                            def do_after_all(self, req, resp, **kwargs):
-                                responder(self, req, resp, **kwargs)
-                                action(req, resp)
+                            do_after_all = _wrap_with_after(
+                                action, responder, resource, True)
 
                             setattr(resource, responder_name, do_after_all)
 
@@ -121,12 +116,108 @@ def do_after_all(self, req, resp, **kwargs):
 
         else:
             responder = responder_or_resource
-
-            @wraps(responder)
-            def do_after_one(self, req, resp, **kwargs):
-                responder(self, req, resp, **kwargs)
-                action(req, resp)
+            do_after_one = _wrap_with_after(action, responder, None, True)
 
             return do_after_one
 
     return _after
+
+# -----------------------------------------------------------------------------
+# Helpers
+# -----------------------------------------------------------------------------
+
+
+def _wrap_with_after(action, responder, resource, is_method=False):
+    """Execute the given action function after a bound responder.
+
+    Args:
+        action: A function with a signature similar to a resource responder
+            method, taking (req, resp).
+        responder: The bound responder to wrap.
+        resource: The resource affected by action.
+        is_method: Is wrapped responder a class method?
+
+    """
+    # NOTE(swistakm): introspect action function do guess if it can handle
+    # additionalresource argument without breaking backwards compatibility
+    spec = inspect.getargspec(action)
+
+    # NOTE(swistakm): create hook before checking what will be actually
+    # decorated. This helps to avoid excessive nesting
+    if len(spec.args) > 2:
+        @wraps(action)
+        def hook(req, resp, resource):
+            action(req, resp, resource)
+    else:
+        @wraps(action)
+        def hook(req, resp, resource):
+            action(req, resp)
+
+    # NOTE(swistakm): method must be decorated differently than normal function
+    if is_method:
+        @wraps(responder)
+        def do_after(self, req, resp, **kwargs):
+            responder(self, req, resp, **kwargs)
+            hook(req, resp, self)
+    else:
+        @wraps(responder)
+        def do_after(req, resp, **kwargs):
+            responder(req, resp, **kwargs)
+            hook(req, resp, resource)
+
+    return do_after
+
+
+def _wrap_with_before(action, responder, resource, is_method=False):
+    """Execute the given action function before a bound responder.
+
+    Args:
+        action: A function with a similar signature to a resource responder
+            method, taking (req, resp, params).
+        responder: The bound responder to wrap.
+        resource: The resource affected by action.
+        is_method: Is wrapped responder a class method?
+
+    """
+    # NOTE(swistakm): introspect action function do guess if it can handle
+    # additional resource argument without breaking backwards compatibility
+    spec = inspect.getargspec(action)
+
+    # NOTE(swistakm): create hook before checking what will be actually
+    # decorated. This allows to avoid excessive nesting
+    if len(spec.args) > 3:
+        @wraps(action)
+        def hook(req, resp, resource, kwargs):
+            action(req, resp, resource, kwargs)
+    else:
+        @wraps(action)
+        def hook(req, resp, resource, kwargs):
+            action(req, resp, kwargs)
+
+    # NOTE(swistakm): method must be decorated differently than normal function
+    if is_method:
+        @wraps(responder)
+        def do_before(self, req, resp, **kwargs):
+            hook(req, resp, self, kwargs)
+            responder(self, req, resp, **kwargs)
+    else:
+        @wraps(responder)
+        def do_before(req, resp, **kwargs):
+            hook(req, resp, resource, kwargs)
+            responder(req, resp, **kwargs)
+
+    return do_before
+
+
+def _wrap_with_hooks(before, after, responder, resource):
+    if after is not None:
+        for action in after:
+            responder = _wrap_with_after(action, responder, resource)
+
+    if before is not None:
+        # Wrap in reversed order to achieve natural (first...last)
+        # execution order.
+        for action in reversed(before):
+            responder = _wrap_with_before(action, responder, resource)
+
+    return responder

tests/test_after_hooks.py

@@ -20,11 +20,21 @@ def fluffiness(req, resp):
     resp.body = 'fluffy'
 
 
+def resource_aware_fluffiness(req, resp, resource):
+    assert resource
+    fluffiness(req, resp)
+
+
 def cuteness(req, resp):
     if resp.body == 'fluffy':
         resp.body += ' and cute'
 
 
+def resource_aware_cuteness(req, resp, resource):
+    assert resource
+    cuteness(req, resp)
+
+
 def fluffiness_in_the_head(req, resp):
     resp.set_header('X-Fluffiness', 'fluffy')
 
@@ -68,6 +78,29 @@ def on_head(self, req, resp):
         self.resp = resp
 
 
+# NOTE(swistakm): we use both type of hooks (class and method)
+# at once for the sake of simplicity
+@falcon.after(resource_aware_cuteness)
+class ClassResourceWithAwareHooks(object):
+
+    # Test that the decorator skips non-callables
+    on_post = False
+
+    def __init__(self):
+        # Test that the decorator skips non-callables
+        self.on_patch = []
+
+    @falcon.after(resource_aware_fluffiness)
+    def on_get(self, req, resp):
+        self.req = req
+        self.resp = resp
+
+    @falcon.after(resource_aware_fluffiness)
+    def on_head(self, req, resp):
+        self.req = req
+        self.resp = resp
+
+
 class ZooResource(object):
 
     def on_get(self, req, resp):
@@ -89,6 +122,9 @@ def before(self):
         self.wrapped_resource = WrappedClassResource()
         self.api.add_route('/wrapped', self.wrapped_resource)
 
+        self.wrapped_resource_aware = ClassResourceWithAwareHooks()
+        self.api.add_route('/wrapped_aware', self.wrapped_resource_aware)
+
     def test_global_hook(self):
         self.assertRaises(TypeError, falcon.API, None, {})
         self.assertRaises(TypeError, falcon.API, None, 0)
@@ -101,6 +137,18 @@ def test_global_hook(self):
         self.simulate_request(self.test_route)
         self.assertEqual(b'fluffy', zoo_resource.resp.body_encoded)
 
+    def test_global_hook_is_resource_aware(self):
+        self.assertRaises(TypeError, falcon.API, None, {})
+        self.assertRaises(TypeError, falcon.API, None, 0)
+
+        self.api = falcon.API(after=resource_aware_fluffiness)
+        zoo_resource = ZooResource()
+
+        self.api.add_route(self.test_route, zoo_resource)
+
+        self.simulate_request(self.test_route)
+        self.assertEqual(b'fluffy', zoo_resource.resp.body_encoded)
+
     def test_multiple_global_hook(self):
         self.api = falcon.API(after=[fluffiness, cuteness])
         zoo_resource = ZooResource()
@@ -190,6 +238,28 @@ def test_wrapped_resource(self):
         self.assertEqual(falcon.HTTP_204, self.srmock.status)
         self.assertEqual([], body)
 
+    def test_wrapped_resource_with_hooks_aware_of_resource(self):
+        expected = b'fluffy and cute'
+
+        self.simulate_request('/wrapped_aware')
+        self.assertEqual(falcon.HTTP_200, self.srmock.status)
+        self.assertEqual(
+            expected, self.wrapped_resource_aware.resp.body_encoded)
+
+        self.simulate_request('/wrapped_aware', method='HEAD')
+        self.assertEqual(falcon.HTTP_200, self.srmock.status)
+
+        self.simulate_request('/wrapped_aware', method='POST')
+        self.assertEqual(falcon.HTTP_405, self.srmock.status)
+
+        self.simulate_request('/wrapped_aware', method='PATCH')
+        self.assertEqual(falcon.HTTP_405, self.srmock.status)
+
+        # decorator does not affect the default on_options
+        body = self.simulate_request('/wrapped_aware', method='OPTIONS')
+        self.assertEqual(falcon.HTTP_204, self.srmock.status)
+        self.assertEqual([], body)
+
     def test_customized_options(self):
         self.api = falcon.API(after=fluffiness)