编写将装饰器应用于所有方法的类装饰器

问题描述

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

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

这几乎可行，但是 @classmethod S需要特殊处理：

It almost works, but @classmethodS 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)

有没有一种方法可以很好地处理此问题？我检查了 @classmethod 装饰的方法，但看不到任何可将它们与其他类型方法区分开的东西。

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.

更新

这是记录的完整解决方案（使用描述符处理 @staticmethod S和 @classmethod S很好，以及aix的技巧来检测 @classmethod S VS普通方法）：

Here is the complete solution for the record (using descriptors to handle @staticmethodS and @classmethodS nicely, and aix's trick to detect @classmethodS 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()

推荐答案

inspect.isclass（meth.im_self）应该告诉您 meth 是否是一个类方法：

inspect.isclass(meth.im_self) should tell you whether meth is a class method:

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

这篇关于编写将装饰器应用于所有方法的类装饰器的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！