问题编写一个将装饰器应用于所有方法的类装饰器

我正在尝试编写一个类装饰器，它将装饰器应用于所有类的方法：

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 装饰方法，但我没有看到任何区别于其他“类型”方法的东西。

更新

这是记录的完整解决方案（使用描述符来处理 @staticmethodS和 @classmethod很好，和艾克斯的技巧来检测 @classmethodS 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()

4792

2017-07-14 15:33

起源

您的DecoratedClassMethod和DecoratedMethod类完全相同。请编辑以提供正确的解决方案。 - Saša Šijak

它们是不同的：DecoratedMethod传递对象实例，而DecoratedClassMethod则不传递。 - Luper Rouch

这些非常相似，必须能够将它们结合起来以避免重复。想一想 self.func(cls or obj, *args, **kwargs)。我知道这不一样，但很简单 if 使用正确测试的声明最终将使您免于拥有这两个几乎完全相同的类。 - robru

答案:

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

2017-07-14 15:47

+1你打败了我。 - SingleNegationElimination

它成功了，谢谢！我用完整的解决方案更新了我的问题。 - Luper Rouch

（评论太长）

我冒昧地添加了指定应该为您的解决方案装饰哪些方法的能力：

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'

2017-08-07 16:39

以上答案不直接适用于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 使用特定属性的类，我更喜欢创建一个类属性。

2018-01-29 19:21

问题 编写一个将装饰器应用于所有方法的类装饰器

答案:

热门问题

问题编写一个将装饰器应用于所有方法的类装饰器