How to create a Python class decorator that is able to wrap instance, class and static methods?

Question

I'd like to create a Python class decorator (*) that would be able to seamlessly wrap all method types the class might have: instance, class and static.

This is the code I have for now, with the parts that break it commented:

def wrapItUp(method):
    def wrapped(*args, **kwargs):
        print "This method call was wrapped!"
        return method(*args, **kwargs)
    return wrapped

dundersICareAbout = ["__init__", "__str__", "__repr__"]#, "__new__"]

def doICareAboutThisOne(cls, methodName):
    return (callable(getattr(cls, methodName))
            and (not (methodName.startswith("__") and methodName.endswith("__"))
            or methodName in dundersICareAbout))

def classDeco(cls):
    myCallables = ((aname, getattr(cls, aname)) for aname in dir(cls) if doICareAboutThisOne(cls, aname))
    for name, call in myCallables:
        print "*** Decorating: %s.%s(...)" % (cls.__name__, name)
        setattr(cls, name, wrapItUp(call))
    return cls

@classDeco
class SomeClass(object):

    def instanceMethod(self, p):
        print "instanceMethod: p =", p

    @classmethod
    def classMethod(cls, p):
        print "classMethod: p =", p

    @staticmethod
    def staticMethod(p):
        print "staticMethod: p =", p


instance = SomeClass()
instance.instanceMethod(1)
#SomeClass.classMethod(2)
#instance.classMethod(2)
#SomeClass.staticMethod(3)
#instance.staticMethod(3)

I'm having two issues trying to make this work:

• When iterating over all callables, how do I find out if it is of an instance, class or static type?
• How to I overwrite the method with a proper wrapped version of it that is invoked correctly for each of those cases?

Currently, this code generates different TypeErrors depending on what commented snippet is uncommented, like:

• TypeError: unbound method wrapped() must be called with SomeClass instance as first argument (got int instance instead)
• TypeError: classMethod() takes exactly 2 arguments (3 given)

---

(*): The same problem is much simpler if you're decorating the methods directly.

Answer 1

Because methods are wrappers for functions, to apply a decorator to a method on a class after the class has been constructed, you have to:

1. Extract the underlying function from the method using its im_func attribute.
2. Decorate the function.
3. Re-apply the wrapper.
4. Overwrite the attribute with the wrapped, decorated function.

It is difficult to distinguish a classmethod from a regular method once the @classmethod decorator has been applied; both kinds of methods are of type instancemethod. However, you can check the im_self attribute and see whether it is None. If so, it's a regular instance method; otherwise it's a classmethod.

Static methods are simple functions (the @staticmethod decorator merely prevents the usual method wrapper from being applied). So you don't have to do anything special for these, it looks like.

So basically your algorithm looks like this:

1. Get the attribute.
2. Is it callable? If not, proceed to the next attribute.
3. Is its type types.MethodType? If so, it is either a class method or an instance method.
   • If its im_self is None, it is an instance method. Extract the underlying function via the im_func attribute, decorate that, and re-apply the instance method: meth = types.MethodType(func, None, cls)
   • If its im_self is not None, it is a class method. Exctract the underlying function via im_func and decorate that. Now you have to reapply the classmethod decorator but you can't because classmethod() doesn't take a class, so there's no way to specify what class it will be attached to. Instead you have to use the instance method decorator: meth = types.MethodType(func, cls, type). Note that the type here is the actual built-in, type.
4. If its type is not types.MethodType then it is a static method or other non-bound callable, so just decorate it.
5. Set the new attribute back onto the class.

These change somewhat in Python 3 -- unbound methods are functions there, IIRC. In any case this will probably need to be completely rethought there.

Answer 2

There is an undocumented function, inspect.classify_class_attrs, which can tell you which attributes are classmethods or staticmethods. Under the hood, it uses isinstance(obj, staticmethod) and isinstance(obj, classmethod) to classify static and class methods. Following that pattern, this works in both Python2 and Python3:

def wrapItUp(method,kind='method'):
    if kind=='static method':
        @staticmethod
        def wrapped(*args, **kwargs):
            return _wrapped(*args,**kwargs)
    elif kind=='class method':
        @classmethod
        def wrapped(cls,*args, **kwargs):
            return _wrapped(*args,**kwargs)                
    else:
        def wrapped(self,*args, **kwargs):
            return _wrapped(self,*args,**kwargs)                                
    def _wrapped(*args, **kwargs):
        print("This method call was wrapped!")
        return method(*args, **kwargs)
    return wrapped
def classDeco(cls):
    for name in (name
                 for name in dir(cls)
                 if (callable(getattr(cls,name))
                     and (not (name.startswith('__') and name.endswith('__'))
                          or name in '__init__ __str__ __repr__'.split()))
                 ):
        method = getattr(cls, name)
        obj = cls.__dict__[name] if name in cls.__dict__ else method
        if isinstance(obj, staticmethod):
            kind = "static method"
        elif isinstance(obj, classmethod):
            kind = "class method"
        else:
            kind = "method"
        print("*** Decorating: {t} {c}.{n}".format(
            t=kind,c=cls.__name__,n=name))
        setattr(cls, name, wrapItUp(method,kind))
    return cls

@classDeco
class SomeClass(object):
    def instanceMethod(self, p):
        print("instanceMethod: p = {}".format(p))
    @classmethod
    def classMethod(cls, p):
        print("classMethod: p = {}".format(p))
    @staticmethod
    def staticMethod(p):
        print("staticMethod: p = {}".format(p))

instance = SomeClass()
instance.instanceMethod(1)
SomeClass.classMethod(2)
instance.classMethod(2)
SomeClass.staticMethod(3)
instance.staticMethod(3)