2011-07-14 20 views
11

我想寫一個適用於裝飾所有的類的方法的類裝飾類裝飾:編寫適用一個裝飾所有方法

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

它幾乎工作,但@classmethod的需要一特殊處理:

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

有沒有辦法很好地處理這個問題?我檢查了@classmethod裝飾方法,但我沒有看到任何東西來區分他們與其他「類型」的方法。

更新

下面是記錄的完整解決方案(用描述符來處理@staticmethod S和@classmethod的很好,和AIX的伎倆來檢測@classmethod s和普通的方法):

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

您的DecoratedClassMethod和DecoratedMethod類完全相同。請編輯正確的解決方案。 –

+0

它們不同:DecoratedMethod傳遞對象實例,而DecoratedClassMethod不傳遞對象實例。 –

+1

這些都非常相似,所以必須將它們結合起來才能避免重複。按照'self.func(cls或obj,* args,** kwargs)'的思路思考。我知道這是不一樣的,但是具有正確測試的簡單的「if」語句最終可以讓你避免使用這兩個幾乎相同的類。 – robru

回答

11

inspect.isclass(meth.im_self)應該告訴你meth是否是一種分類方法:

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 
+1

+1你擊敗了我。 – SingleNegationElimination

+0

它做的伎倆,謝謝!我用完整的解決方案更新了我的問題。 –

1

(太多了評論長)

我把添加指定哪些方法應該得到裝飾解決方案能力的自由:

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 

用法:

@class_decorator('some_method') 
class Foo(object): 

    def some_method(self): 
     print 'I am decorated' 

    def another_method(self): 
     print 'I am NOT decorated' 
0

以上答案不直接適用於python3。基於其他偉大的答案我已經能夠想出以下解決方案:

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 

現在,只要在元組methods_to_override的方法被調用時,髒標誌設置。當然,其他任何東西都可以放在那裏。沒有必要在其方法需要被覆蓋的類中包含DecorateMethod類。但是,由於DecorateMehod對類使用了特定的屬性,因此我傾向於創建類屬性。