Escribir un decorador de clase que aplique un decorador a todos los métodos.

Estoy tratando de escribir un decorador de clase que aplique un decorador a todos los métodos de la clase:

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() 

Casi funciona, pero @classmethod S necesita un tratamiento especial:

 $ python test.py before __init__ after before 10 after baz baz before Traceback (most recent call last): File "test.py", line 44, in  obj.bar() File "test.py", line 7, in wrapper ret = func(*args, **kwargs) TypeError: bar() takes exactly 1 argument (2 given) 

¿Hay alguna manera de manejar bien este problema? Inspeccioné métodos de decoración de @classmethod , pero no veo nada para diferenciarlos de otros “tipos” de métodos.

Actualizar

Aquí está la solución completa para el registro (usando descriptores para manejar @staticmethod S y @classmethod S, y el truco de aix para detectar los métodos normales de @classmethod S VS):

 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() 

inspect.isclass(meth.im_self) debe decirle si la meth es un método de clase:

 def decorate_class(cls): for name, meth in inspect.getmembers(cls, inspect.ismethod): if inspect.isclass(meth.im_self): print '%s is a class method' % name # TODO ... return cls 

(Demasiado largo para un comentario)

Me tomé la libertad de agregar la capacidad de especificar qué métodos deberían decorar su solución:

 def class_decorator(*method_names): def wrapper(cls): for name, meth in inspect.getmembers(cls): if name in method_names or len(method_names) == 0: if inspect.ismethod(meth): if inspect.isclass(meth.im_self): # meth is a classmethod setattr(cls, name, VerifyTokenMethod(meth)) else: # meth is a regular method setattr(cls, name, VerifyTokenMethod(meth)) elif inspect.isfunction(meth): # meth is a staticmethod setattr(cls, name, VerifyTokenMethod(meth)) return cls return wrapper 

Uso:

 @class_decorator('some_method') class Foo(object): def some_method(self): print 'I am decorated' def another_method(self): print 'I am NOT decorated' 

Las respuestas anteriores no se aplican directamente a python3. Basándome en las otras excelentes respuestas, he podido encontrar la siguiente solución:

 import inspect import types import networkx as nx def override_methods(cls): for name, meth in inspect.getmembers(cls): if name in cls.methods_to_override: setattr(cls, name, cls.DecorateMethod(meth)) return cls @override_methods class DiGraph(nx.DiGraph): methods_to_override = ("add_node", "remove_edge", "add_edge") class DecorateMethod: def __init__(self, func): self.func = func def __get__(self, obj, cls=None): def wrapper(*args, **kwargs): ret = self.func(obj, *args, **kwargs) obj._dirty = True # This is the attribute I want to update return ret return wrapper def __init__(self): super().__init__() self._dirty = True 

Ahora, cada methods_to_override se llama a un método en la tupla methods_to_override , se establece el indicador sucio. Por supuesto, cualquier otra cosa se puede poner allí también. No es necesario incluir la clase DecorateMethod en la clase cuyos métodos deben anularse. Sin embargo, como DecorateMehod usa atributos específicos para la clase, prefiero hacer un atributo de clase.