Search code examples
pythondecoratorclass-method

How to decorate class or static methods


I am writing a generic class decorator which needs to apply a decorator to each method. My first approach is something like this:

def class_decorator(cls):
    for name, member in vars(cls).items():
        # Ignore anything that is not a method
        if not isinstance(member, (types.FunctionType, types.BuiltinFunctionType, classmethod, staticmethod)):
            continue

        setattr(cls, name, method_decorator(member))

    return cls

The decorator itself is not very important. Looks something like this:

def method_decorator(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        # do something
        return fn(*args, **kwargs):

    return wrapper

Once I tested it, I ran into the problem that this does not work with static or class methods, and the following error is raised from functools.wraps:

AttributeError: 'classmethod' object has no attribute '__module__'

Yeah, classmethod or staticmethods are not normal functions, not even callables. Generally if you need to decorate a classmethod, you first apply your decorator and then the classmethod decorator, but since this is a class decorator, I cannot influence the order of the decorators.

Any good solution for this?


Solution

  • After playing around for a while, I have found a solution that looks to me better than other approaches in SO. Maybe this can be helpful to somebody.

    Basically the idea is the following:

    • Detect members which are class or static methods
    • Get the function object wrapped within these methods
    • Apply the decorator to this function
    • Wrap the decorated function within a classmethod or staticmethod instance
    • Store it in the class again

    The code looks like this:

    def class_decorator(cls):
        for name, member in vars(cls).items():
            # Good old function object, just decorate it
            if isinstance(member, (types.FunctionType, types.BuiltinFunctionType)):
                setattr(cls, name, method_decorator(member))
                continue
    
            # Static and class methods: do the dark magic
            if isinstance(member, (classmethod, staticmethod)):
                inner_func = member.__func__
                method_type = type(member)
                decorated = method_type(method_decorator(inner_func))
                setattr(cls, name, decorated)
                continue
    
            # We don't care about anything else
    
        return cls