Search code examples
pythonfunctionmethodsargumentsdecorator

Using the same decorator (with arguments) with functions and methods


I have been trying to create a decorator that can be used with both functions and methods in python. This on it's own is not that hard, but when creating a decorator that takes arguments, it seems to be.

class methods(object):
    def __init__(self, *_methods):
        self.methods = _methods

    def __call__(self, func): 
        def inner(request, *args, **kwargs):
            print request
            return func(request, *args, **kwargs)
        return inner

    def __get__(self, obj, type=None):
        if obj is None:
            return self
        new_func = self.func.__get__(obj, type)
        return self.__class__(new_func)

The above code wraps the function/method correctly, but in the case of a method, the request argument is the instance it is operating on, not the first non-self argument.

Is there a way to tell if the decorator is being applied to a function instead of a method, and deal accordingly?


Solution

  • To expand on the __get__ approach. This can be generalized into a decorator decorator.

    class _MethodDecoratorAdaptor(object):
        def __init__(self, decorator, func):
            self.decorator = decorator
            self.func = func
        def __call__(self, *args, **kwargs):
            return self.decorator(self.func)(*args, **kwargs)
        def __get__(self, instance, owner):
            return self.decorator(self.func.__get__(instance, owner))
    
    def auto_adapt_to_methods(decorator):
        """Allows you to use the same decorator on methods and functions,
        hiding the self argument from the decorator."""
        def adapt(func):
            return _MethodDecoratorAdaptor(decorator, func)
        return adapt
    

    In this way you can just make your decorator automatically adapt to the conditions it is used in.

    def allowed(*allowed_methods):
        @auto_adapt_to_methods
        def wrapper(func):
            def wrapped(request):
                if request not in allowed_methods:
                    raise ValueError("Invalid method %s" % request)
                return func(request)
            return wrapped
        return wrapper
    

    Notice that the wrapper function is called on all function calls, so don't do anything expensive there.

    Usage of the decorator:

    class Foo(object):
        @allowed('GET', 'POST')
        def do(self, request):
            print "Request %s on %s" % (request, self)
    
    @allowed('GET')
    def do(request):
        print "Plain request %s" % request
    
    Foo().do('GET')  # Works
    Foo().do('POST') # Raises