Escribir un decorador de clase que se aplica un decorador a todos los métodos

Question

Estoy tratando de escribir un decorador de clase que se aplica a un decorador para 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 necesitan un @classmethod S 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 <module> 
    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 este problema? Inspeccioné @classmethod métodos decorados, pero no veo nada para diferenciarlos de otros "tipos" de métodos.

actualización

Aquí es la solución completa para el registro (usando descriptores de manejar @staticmethod S y @classmethod S agradable, y el truco de Aix para detectar @classmethod métodos normales 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() 

Answer 1

inspect.isclass(meth.im_self) debería decirle si 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 

Answer 2

(demasiado tiempo para un comentario)

he tomado la libertad de añadir la capacidad de especificar qué métodos deben quedar decoradas para 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' 

Answer 3

Lo anterior las respuestas no se aplican directamente a python3. Sobre la base de los otros grandes respuestas que he sido capaz de llegar a 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 vez que un método en la tupla methods_to_override se llama, la bandera sucia se establece. 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 sustituirse. Sin embargo, como DecorateMehod usa atributos específicos para la clase, prefiero hacer un atributo de clase.