Class method decorator with self arguments?

Viewed 203065

How do I pass a class field to a decorator on a class method as an argument? What I want to do is something like:

class Client(object):
    def __init__(self, url):
        self.url = url

    @check_authorization("some_attr", self.url)
    def get(self):
        do_work()

It complains that self does not exist for passing self.url to the decorator. Is there a way around this?

8 Answers

Another option would be to abandon the syntactic sugar and decorate in the __init__ of the class.

def countdown(number):
    def countdown_decorator(func):
        def func_wrapper():
            for index in reversed(range(1, number+1)):
                print(index)
            func()
        return func_wrapper
    return countdown_decorator

class MySuperClass():
    def __init__(self, number):
        self.number = number
        self.do_thing = countdown(number)(self.do_thing)
    
    def do_thing(self):
        print('im doing stuff!')


myclass = MySuperClass(3)

myclass.do_thing()

which would print

3
2
1
im doing stuff!

I know this issue is quite old, but the below workaround hasn't been proposed before. The problem here is that you can't access self in a class block, but you can in a class method.

Let's create a dummy decorator to repeat a function some times.

import functools
def repeat(num_rep):
    def decorator_repeat(func):
        @functools.wraps(func)
        def wrapper_repeat(*args, **kwargs):
            for _ in range(num_rep):
                value = func(*args, **kwargs)
            return 
        return wrapper_repeat
    return decorator_repeat
class A:
    def __init__(self, times, name):
        self.times = times
        self.name = name
    
    def get_name(self):
        @repeat(num_rep=self.times)
        def _get_name():
            print(f'Hi {self.name}')
        _get_name()

I know this is an old question, but this solution has not been mentioned yet, hopefully it may help someone even today, after 8 years.

So, what about wrapping a wrapper? Let's assume one cannot change the decorator neither decorate those methods in init (they may be @property decorated or whatever). There is always a possibility to create custom, class-specific decorator that will capture self and subsequently call the original decorator, passing runtime attribute to it.

Here is a working example (f-strings require python 3.6):

import functools

# imagine this is at some different place and cannot be changed
def check_authorization(some_attr, url):
        def decorator(func):
                @functools.wraps(func)
                def wrapper(*args, **kwargs):
                        print(f"checking authorization for '{url}'...")
                        return func(*args, **kwargs)
                return wrapper
        return decorator

# another dummy function to make the example work
def do_work():
        print("work is done...")

###################
# wrapped wrapper #
###################
def custom_check_authorization(some_attr):
        def decorator(func):
                # assuming this will be used only on this particular class
                @functools.wraps(func)
                def wrapper(self, *args, **kwargs):
                        # get url
                        url = self.url
                        # decorate function with original decorator, pass url
                        return check_authorization(some_attr, url)(func)(self, *args, **kwargs)
                return wrapper
        return decorator
        
#############################
# original example, updated #
#############################
class Client(object):
        def __init__(self, url):
                self.url = url
    
        @custom_check_authorization("some_attr")
        def get(self):
                do_work()

# create object
client = Client(r"https://stackoverflow.com/questions/11731136/class-method-decorator-with-self-arguments")

# call decorated function
client.get()

output:

checking authorisation for 'https://stackoverflow.com/questions/11731136/class-method-decorator-with-self-arguments'...
work is done...

It will be very useful to have a general-purpose utility, that can turn any decorator for functions, into decorator for methods. I thought about it for an hour, and actually come up with one:

from typing import Callable
Decorator = Callable[[Callable], Callable]

def decorate_method(dec_for_function: Decorator) -> Decorator:

    def dec_for_method(unbounded_method) -> Callable:
        # here, `unbounded_method` will be a unbounded function, whose
        # invokation must have its first arg as a valid `self`. When it 
        # return, it also must return an unbounded method.
        def decorated_unbounded_method(self, *args, **kwargs):
            @dec_for_function
            def bounded_method(*args, **kwargs):
                return unbounded_method(self, *args, **kwargs)
            return bounded_method(*args, **kwargs)

        return decorated_unbounded_method

    return dec_for_method

The usage is:

# for any decorator (with or without arguments)
@some_decorator_with_arguments(1, 2, 3)
def xyz(...): ...

# use it on a method:
class ABC:
  @decorate_method(some_decorator_with_arguments(1, 2, 3))
  def xyz(self, ...): ...

Test:

def dec_for_add(fn):
    """This decorator expects a function: (x,y) -> int.

    If you use it on a method (self, x, y) -> int, it will fail at runtime.
    """
    print(f"decorating: {fn}")
    def add_fn(x,y):
        print(f"Adding {x} + {y} by using {fn}")
        return fn(x,y)
    return add_fn


@dec_for_add
def add(x,y):
    return x+y

add(1,2)  # OK!


class A:
    @dec_for_add
    def f(self, x, y):
        # ensure `self` is still a valid instance
        assert isinstance(self, A)
        return x+y

# TypeError: add_fn() takes 2 positional arguments but 3 were given
# A().f(1,2)
    

class A:
    @decorate_method(dec_for_add)
    def f(self, x, y):
        # ensure `self` is still a valid instance
        assert isinstance(self, A)
        return x+y

# Now works!!
A().f(1,2)
Related