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

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

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

回顾

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

函数装饰器

函数形式

不再赘述,直接看代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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 的函数,这个函数对原函数进行修饰。

调用代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
@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))

输出结果

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

类形式

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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

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

类装饰器

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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__ 等于执行了原类的构造函数。能不能再给力一点呢?如果我们想在修饰类的同时顺便把他的成员函数也修饰了呢?

当然可以

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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 函数,但原类的函数名不应该暴露给装饰器,如果原类是一个黑盒呢?能不能再给力一点呢?

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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,接下来返回一个 wrappertarget 进行包装即可。

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

最终的代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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 判断目标属性是否是函数。至于其余剩下的改动,就留给读者思考了。

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

评论