Write a decorator to apply another decorator with arguments for all methods of a class

Question

Coming from this post. The accepted answer works nicely with decorators which take no argument. I'm trying to extend this solution to make it take arguments for the applying decorator.

In detail, I have functions which make external api calls. Because those calls fail frequently, I applied retry decorator from this library to all the functions. To avoid putting the @retry(...) line again and again for all the functions, I decided to centralize them in a class. I created RetryClass and put all the functions as classmethod in the class. Now, I'm looking for a way to apply the retry decorator for all the methods of the class, so that I can just keep adding new methods in the class and it will automatically apply the retry decorator for the new methods.

Note: the retry decorator takes arguments.

@retry(wait_random_min=100, wait_random_max=300, stop_max_attempt_number=3)

Here is my code:

from retrying import retry


def for_all_methods(decorator):
    def decorate(cls):
        for attr in cls.__dict__:
            if callable(getattr(cls, attr)):
                setattr(cls, attr, decorator(getattr(cls, attr)))
        return cls
    return decorate


@for_all_methods(retry(wait_random_min=100, wait_random_max=300, stop_max_attempt_number=3))
class RetryClass(object):

    @classmethod
    def a(cls):
        pass


def test():
    RetryClass.a()
    return

This throws the following error:

Traceback (most recent call last):
  File "/Applications/PyCharm.app/Contents/helpers/pydev/pydevd.py", line 1596, in <module>
    globals = debugger.run(setup['file'], None, None, is_module)
  File "/Applications/PyCharm.app/Contents/helpers/pydev/pydevd.py", line 974, in run
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "/Users/gyoho/Datatron/Dev/class-decorator/main.py", line 26, in <module>
    test()
  File "/Users/gyoho/Datatron/Dev/class-decorator/main.py", line 22, in test
    RetryClass.a()
TypeError: unbound method a() must be called with RetryClass instance as first argument (got nothing instead)

However, commenting out the class decorator runs with no error. Is there anything I'm missing?

When you want to apply it to all methods, you may be better off with a [meta class](https://stackoverflow.com/questions/100003/what-is-a-metaclass-in-python#6581949). — mdh, Feb 26 '17 at 22:04

score 1 · Answer 1 · answered Feb 26 '17 at 22:05

The problem is that @classmethod is no longer the top level decorator for a(). RetryClass.a is currently decorated with both @classmethod and @retry, in that order. RetryClass is equivalent to:

class RetryClass(object):

    @retry(wait_random_min=100, wait_random_max=300, stop_max_attempt_number=3)
    @classmethod
    def a(cls):
        pass

Your class needs to be equivalent to:

class RetryClass(object):

    @classmethod
    @retry(wait_random_min=100, wait_random_max=300, stop_max_attempt_number=3)
    def a(cls):
        pass

I guess I just need to avoid the `classmethod` or `staticmethod` then. Instantiate this class and use its method? — gyoho, Feb 27 '17 at 03:58

score 0 · Accepted Answer · answered Feb 26 '17 at 22:29

0

classmethod and staticmethod must be the last decorator on a method, because they return descriptors and not a function. (Which are trickier to decorate). You could detect if a method is already a classmethod or staticmethod, and then your decorate function would look somewhat like this:

def decorate(cls):
    for attr in cls.__dict__:
        possible_method = getattr(cls, attr)
        if not callable(possible_method):
            continue

        if not hasattr(possible_method, "__self__"):
            raw_function = cls.__dict__[attr].__func__
            decorated_method = decorator(raw_function)
            decorated_method = staticmethod(decorated_method)

        if type(possible_method.__self__) == type:
            raw_function = cls.__dict__[attr].__func__
            decorated_method = decorator(raw_function)
            decorated_method = classmethod(decorated_method)

        elif possible_method.__self__ is None:
            decorated_method = decorator(possible_method)

        setattr(cls, attr, decorated_method)

    return cls

answered Feb 26 '17 at 22:29

Dean Fenster

2,345
1
18
27

Maybe, it'll be easier just to instantiate an object from the class and use its methods then? – gyoho Feb 27 '17 at 04:58
@gyoho It would possibly be simpler, bur keep in mind you'd still have trouble with `staticmethod` and `property` – Dean Fenster Feb 27 '17 at 07:51
I removed the `@classmethod` and changed `cls` to `self`, but didn't put `@staticmethod`. It runs fine this time. What could be the drawback for doing this? – gyoho Feb 27 '17 at 07:55
The main drawback as I see it is that you might want that functionality in the furture. Also, depending on what your methods do, you might not want to have multiple instances (e.g. multiple database connection) – Dean Fenster Feb 27 '17 at 07:58
I instantiated an object in `__init__.py` and import the object as necessary, so it's like a singleton object. Because the class has no variable, I guess for this particular case it's OK to do this way? – gyoho Feb 27 '17 at 08:03

Write a decorator to apply another decorator with arguments for all methods of a class

2 Answers2