Ticket #13879: ticket13879-r16259.diff

tests/regressiontests/decorators/tests.py

@@ -3 +3 @@
 from django.contrib.auth.decorators import login_required, permission_required, user_passes_test
 from django.contrib.admin.views.decorators import staff_member_required
 from django.http import HttpResponse, HttpRequest, HttpResponseNotAllowed
-from django.utils.decorators import method_decorator
+from django.utils.decorators import method_decorator, method_decorator_with_args
 from django.utils.functional import allow_lazy, lazy, memoize
 from django.utils.unittest import TestCase
 from django.views.decorators.http import require_http_methods, require_GET, require_POST, require_safe
@@ -201 +201 @@
 
         self.assertEqual(Test.method.__doc__, 'A method')
         self.assertEqual(Test.method.im_func.__name__, 'method')
+
+
+# For testing method_decorator_with_args, a decorator factory, that returns
+# a decorator that assumes a single argument. We will get type arguments if
+# there is a mismatch in the number of arguments.
+def prefix_arg(prefix):
+    def dec(func):
+        def wrapper(arg):
+            return func(prefix + arg)
+        return wraps(func)(wrapper)
+    return dec
+
+prefix_arg_m = method_decorator_with_args(prefix_arg)
+
+# For testing method_decorator_with_args, a decorator factory that returns
+# a decorator that adds an attribute to the function.
+def add_attr(attname):
+    def dec(func):
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+        setattr(wrapper, attname, True)
+        return wraps(func)(wrapper)
+    return dec
+
+add_attr_m = method_decorator_with_args(add_attr)
+
+class MethodDecoratorWithArgsTests(TestCase):
+    """
+    Tests for method_decorator_with_args
+    """
+    def test_preserve_signature(self):
+        class Test(object):
+            @prefix_arg_m("test:")
+            def say(self, arg):
+                return arg
+
+        self.assertEqual("test:hello", Test().say("hello"))
+
+    def test_preserve_attributes(self):
+        # Sanity check add_attr
+        @add_attr('myattr')
+        @add_attr('myattr2')
+        def func():
+            pass
+
+        self.assertEqual(getattr(func, 'myattr', False), True)
+        self.assertEqual(getattr(func, 'myattr2', False), True)
+
+        # Now check method_decorator
+        class Test(object):
+            @add_attr_m('myattr')
+            @add_attr_m('myattr2')
+            def method(self):
+                "A method"
+                pass
+
+        self.assertEqual(getattr(Test().method, 'myattr', False), True)
+        self.assertEqual(getattr(Test().method, 'myattr2', False), True)
+
+        self.assertEqual(getattr(Test.method, 'myattr', False), True)
+        self.assertEqual(getattr(Test.method, 'myattr2', False), True)
+
+        self.assertEqual(Test.method.__doc__, 'A method')
+        self.assertEqual(Test.method.im_func.__name__, 'method')

django/utils/decorators.py

@@ -40 +40 @@
     return _dec
 
 
+def method_decorator_with_args(decorator_factory):
+    """
+    Converts a function decorator factory (e.g. `permission_required`) into a
+    method decorator factory.
+    """
+    def _method_dec_factory(*args, **kwargs):
+        return method_decorator(decorator_factory(*args, **kwargs))
+    update_wrapper(_method_dec_factory, decorator_factory)
+    # Change the name to aid debugging.
+    _method_dec_factory.__name__ = 'method_decorator_with_args(%s)' % decorator_factory.__name__
+    return _method_dec_factory
+
+
 def decorator_from_middleware_with_args(middleware_class):
     """
     Like decorator_from_middleware, but returns a function