编写一个装饰类的装饰器类

快速编写一个装饰类的装饰器类 (a decorator class to decorate a class)，调用方便、扩展性强、内附代码实现

在我们会编写函数装饰器用于装饰函数、类装饰器用于装饰函数后，我们很自然会想到一个问题，能否编写类装饰器装饰一个类？我们能否通过仅仅对类装饰，却能 Hook 掉这个类的所有成员函数以达到方便扩展的目的？本文将快速回顾前几种装饰器，并最终得到一个装饰类的全能装饰器类。

回顾

我们先来回顾一下前几种装饰器的基本形式，并且假设我们现在需要记录函数执行的日志、包括函数的输入和输出

函数装饰器

函数形式

不再赘述，直接看代码

from functools import wraps


def Logger(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print('call %s() with args: %s, kwargs: %s' % (func.__name__, args, kwargs))
        ret = func(*args, **kwargs)
        print('%s() return %s' % (func.__name__, ret))
        return ret

    return wrapper


def NamedLogger(name):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            print('%s: call %s() with args: %s, kwargs: %s' % (name, func.__name__, args, kwargs))
            ret = func(*args, **kwargs)
            print('%s: %s() return %s' % (name, func.__name__, ret))
            return ret

        return wrapper

    return decorator

这里我们编写了两个函数用作装饰器。对于无参数的装饰器，我们直接返回一个 wrapper 对函数进行修饰。而对于无参数的装饰器，我们返回一个返回 wrapper 的函数，这个函数对原函数进行修饰。

调用代码如下

@Logger
def add(a, b):
    return a + b


@NamedLogger('myLogger')
def sub(a, b):
    return a - b


if __name__ == '__main__':
    print(add(1, 2))
    print(sub(1, 2))

输出结果

call add() with args: (1, 2), kwargs: {}
add() return 3
3
myLogger: call sub() with args: (1, 2), kwargs: {}
myLogger: sub() return -1
-1

类形式

我们能否编写一个类用于装饰函数呢，答案是肯定的，上代码

class Logger:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        print(f'call {self.func.__name__}() with args: {args}, kwargs: {kwargs}')
        ret = self.func(*args, **kwargs)
        print(f'{self.func.__name__}() return: {ret}')
        return ret


class NamedLogger:
    def __init__(self, name):
        self.name = name

    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            print(f'{self.name}: call {func.__name__}() with args: {args}, kwargs: {kwargs}')
            ret = func(*args, **kwargs)
            print(f'{self.name}: {func.__name__}() return: {ret}')
            return ret

        return wrapper

注意带参数类 Logger 和不带参数的类 NamedLogger 中 __init__ 的区别，在不带参数的类中，__init__ 函数实际上是把原函数作为参数传入了，而在带参数的类中，__init__ 函数的参数作为装饰器本身。同时两者都实现了 __call__ 函数，用于模拟函数的行为，但其中的实现大不相同。Logger 中初始化已经得到了函数，因此在 __call__ 中直接调用即可，而 NamedLogger 中 __init__ 仅初始化了装饰器本身，因此 __call__ 需要返回一个原函数的包装 wrapper。

使用和函数形式相同的调用代码，输出结果和上述一致。

类装饰器

介绍完了函数装饰器，下面我们来介绍类装饰器。根据前面的经验，既然函数装饰器需要返回一个函数，那么类装饰器我们返回一个类就好了。因此我们可以快速编写这样一段代码

def NamedLogger(name):
    class ClassWrapper:
        def __init__(self, cls):
            self.old_class = cls

        def __call__(self, *args, **kwargs):
            print(f'{name}: call {self.old_class.__name__}() with args: {args}, kwargs: {kwargs}')
            ret = self.old_class(*args, **kwargs)
            print(f'{name}: {self.old_class.__name__}() return: {ret}')
            return self.old_class(*args, **kwargs)

    return ClassWrapper


@NamedLogger('myLogger')
class Adder:
    def __init__(self, a, b):
        self.ret = a + b


if __name__ == '__main__':
    adder = Adder(1, 2)
    print(adder.ret)

注意这里 __call__ 等于执行了原类的构造函数。能不能再给力一点呢？如果我们想在修饰类的同时顺便把他的成员函数也修饰了呢？

当然可以

def NamedLogger(name):
    class ClassWrapper:
        def __init__(self, cls):
            self.old_class = cls

        def __call__(self, *args, **kwargs):
            self.old_object = self.old_class(*args, **kwargs)
            return self

        def add(self, *args, **kwargs):
            print(f'{name}: call {self.old_class.__name__}.add() with args: {args}, kwargs: {kwargs}')
            ret = self.old_object.add(*args, **kwargs)
            print(f'{name}: {self.old_class.__name__}.add() return: {ret}')
            return ret

    return ClassWrapper


@NamedLogger('myLogger')
class Adder:
    def add(self, a, b):
        return a + b


if __name__ == '__main__':
    adder = Adder()
    print(adder.add(1, 2))

我们在这里通过 __call__ 返回了 ClassWrapper 本身，这样在 adder = Adder() 调用时本质是实际 adder 是拿到了 ClassWrapper 这个对象，然后我们通过直接定义一个 add() 函数来接管原类的 add 函数。

但是问题又来了，这个实现直接定义了一个 add 函数，但原类的函数名不应该暴露给装饰器，如果原类是一个黑盒呢？能不能再给力一点呢？

额，有点复杂，不过还是可以做到的

from functools import wraps


def NamedLogger(name):
    class ClassWrapper:
        def __init__(self, cls):
            self.old_class = cls

        def __call__(self, *args, **kwargs):
            self.old_object = self.old_class(*args, **kwargs)
            return self

        def __getattr__(self, func_name):
            target = getattr(self.old_object, func_name)

            @wraps(target)
            def wrapper(*args, **kwargs):
                print(f"{name}: call {func_name}() with args: {args}, kwargs: {kwargs}")
                result = target(*args, **kwargs)
                print(f"{name}: {func_name}()  return: {result}")
                return result

            return wrapper

    return ClassWrapper

我们重载 __getattr__ 函数，然后通过 getattr 得到函数目标 target，接下来返回一个 wrapper 对 target 进行包装即可。

好了，我们差不多完成了类装饰器的雏形，这个装饰器可以对一个类的所有成员函数进行修饰，非常方便。我们最后再补充一点细节，比如我们的类装饰器能否同时兼容有参数形式和无参数形式？如果原类调用了成员变量怎么办？按照现在的实现同样会返回一个 wrapper。

最终的代码如下

from functools import wraps, partial


def NamedLogger(cls=None, name='default Logger'):
    if cls is None:
        return partial(NamedLogger, name=name)

    @wraps(cls, updated=())  # see https://stackoverflow.com/a/65470430/7620214
    class ClassWrapper:
        def __init__(self, *cls_args, **cls_kwargs):
            self.old_object = cls(*cls_args, **cls_kwargs)

        def __getattr__(self, func_name):
            target = getattr(self.old_object, func_name)
            if not callable(target):
                return target

            @wraps(target)
            def wrapper(*args, **kwargs):
                print(f"{name}: call {func_name}() with args: {args}, kwargs: {kwargs}")
                result = target(*args, **kwargs)
                print(f"{name}: {func_name}()  return: {result}")
                return result

            return wrapper

    return ClassWrapper


@NamedLogger(name='myLogger')
class Adder:
    def add(self, a, b):
        return a + b


@NamedLogger
class Adder2:
    def __init__(self):
        self.last = 0

    def add(self, a, b):
        self.last = a + b
        return self.last


if __name__ == '__main__':
    adder = Adder()
    print(adder.add(1, 2))
    adder2 = Adder2()
    print(adder2.add(1, 2))
    print(adder2.last)

代码的实现变得更简洁了一点，__call__ 也被干掉了，并且同时支持带参数和不带参数两种不同形式的装饰方法。此外，我们还用 callable 判断目标属性是否是函数。至于其余剩下的改动，就留给读者思考了。

最后，你可以在 这里 下载到最终的源代码。