I'm trying to write a class decorator that applies a decorator to all the class' methods:
import inspect
def decorate_func(func):
def wrapper(*args, **kwargs):
print "before"
ret = func(*args, **kwargs)
print "after"
return ret
for attr in "__module__", "__name__", "__doc__":
setattr(wrapper, attr, getattr(func, attr))
return wrapper
def decorate_class(cls):
for name, meth in inspect.getmembers(cls, inspect.ismethod):
setattr(cls, name, decorate_func(meth))
return cls
@decorate_class
class MyClass(object):
def __init__(self):
self.a = 10
print "__init__"
def foo(self):
print self.a
@staticmethod
def baz():
print "baz"
@classmethod
def bar(cls):
print "bar"
obj = MyClass()
obj.foo()
obj.baz()
MyClass.baz()
obj.bar()
MyClass.bar()
It almost works, but @classmethod
S need a special treatment:
$ python test.py
before
__init__
after
before
10
after
baz
baz
before
Traceback (most recent call last):
File "test.py", line 44, in <module>
obj.bar()
File "test.py", line 7, in wrapper
ret = func(*args, **kwargs)
TypeError: bar() takes exactly 1 argument (2 given)
Is there a way to handle this problem nicely ? I inspected @classmethod
decorated methods, but I don't see anything to differentiate them from other "types" of methods.
Update
Here is the complete solution for the record (using descriptors to handle @staticmethod
S and @classmethod
S nicely, and aix's trick to detect @classmethod
S VS normal methods):
import inspect
class DecoratedMethod(object):
def __init__(self, func):
self.func = func
def __get__(self, obj, cls=None):
def wrapper(*args, **kwargs):
print "before"
ret = self.func(obj, *args, **kwargs)
print "after"
return ret
for attr in "__module__", "__name__", "__doc__":
setattr(wrapper, attr, getattr(self.func, attr))
return wrapper
class DecoratedClassMethod(object):
def __init__(self, func):
self.func = func
def __get__(self, obj, cls=None):
def wrapper(*args, **kwargs):
print "before"
ret = self.func(*args, **kwargs)
print "after"
return ret
for attr in "__module__", "__name__", "__doc__":
setattr(wrapper, attr, getattr(self.func, attr))
return wrapper
def decorate_class(cls):
for name, meth in inspect.getmembers(cls):
if inspect.ismethod(meth):
if inspect.isclass(meth.im_self):
# meth is a classmethod
setattr(cls, name, DecoratedClassMethod(meth))
else:
# meth is a regular method
setattr(cls, name, DecoratedMethod(meth))
elif inspect.isfunction(meth):
# meth is a staticmethod
setattr(cls, name, DecoratedClassMethod(meth))
return cls
@decorate_class
class MyClass(object):
def __init__(self):
self.a = 10
print "__init__"
def foo(self):
print self.a
@staticmethod
def baz():
print "baz"
@classmethod
def bar(cls):
print "bar"
obj = MyClass()
obj.foo()
obj.baz()
MyClass.baz()
obj.bar()
MyClass.bar()
See Question&Answers more detail:
os