Zwlin's Blog

Python装饰器(二)

2020/01/04

可接受参数的装饰器

假设编写一个为函数添加日志功能的装饰器,但是又允许用户指定日志的等级以及一些其他的细节作为参数,下面是定义这个装饰器的可能做法:

 1from functools import wraps
 2
 3import logging
 4
 5
 6def logged(level, name=None, message=None):
 7    '''
 8    Add logging to a function,if name and message is aren't specified,they default to the function's module and name.
 9    :param level: the logging level
10    :param name:  the logger name
11    :param message: the log message
12    '''
13
14    def decorate(func):
15        logname = name if name else func.__module__
16        log = logging.getLogger(logname)
17        logmsg = message if message else func.__name__
18
19        @wraps(func)
20        def wrapper(*args, **kwargs):
21            log.log(level, logmsg)
22            return func(*args, **kwargs)
23
24        return wrapper
25
26    return decorate
27
28
29# Example use
30@logged(logging.DEBUG)
31def add(x, y):
32    return x + y
33
34
35@logged(logging.CRITICAL, 'example')
36def spam():
37    print("Spam!")

初看上去这个实现很有技巧性,但是其中的思想相对来说是很简单的,最外层的logged()函数接受所需的参数,并让他们对装饰器的内层函数可见,内层的decorate()函数接受一个参数并给他加上一个包装层,关键部分在于这个包装层可以使用传递给logged()的参数。

编写一个可接受参数的装饰器是需要一定技巧的,因为这会涉及底层的调用顺序,具体来说,如果有这样的代码:

1@decorator(x,y,z)
2def func(a,b):
3    pass

装饰的过程会按照下列方式来进行:

1def func(a,b):
2    pass
3func = decorator(x,y,z)(func)

decorator(x,y,z)的结果必须是一个可调用对象,这个对象反过来接受一个函数作为参数输入,并对其进行包装。

属性可由用户修改的装饰器

编写一个装饰器来来包装函数,但是可以让用户调整装饰器的属性,这样运行时能够控制装饰器的行为。接下来给出的解决方案对上一节的示例作了扩展,引入了访问器函数accessor function,通过使用nonlocal关键字声明变量来修改装饰器内部的属性,之后把访问器函数作为函数属性附加到包装函数上。

 1from functools import wraps, partial
 2import logging
 3
 4
 5# Utility decorator to attach a function as an attribute of obj
 6def attach_wrapper(obj, func=None):
 7    if func is None:
 8        return partial(attach_wrapper, obj)
 9    setattr(obj, func.__name__, func)
10    return func
11
12
13def logged(level, name=None, message=None):
14    '''
15    Add logging to a function,if name and message is aren't specified,
16    they default to the function's module and name.
17    :param level: the logging level
18    :param name:  the logger name
19    :param message: the log message
20    '''
21
22    def decorate(func):
23        logname = name if name else func.__module__
24        log = logging.getLogger(logname)
25        logmsg = message if message else func.__name__
26
27        @wraps(func)
28        def wrapper(*args, **kwargs):
29            log.log(level, logmsg)
30            return func(*args, **kwargs)
31
32        @attach_wrapper(wrapper)
33        def set_level(newlevel):
34            nonlocal level
35            level = newlevel
36
37        @attach_wrapper(wrapper)
38        def set_message(newmsg):
39            nonlocal logmsg
40            logmsg = newmsg
41
42        return wrapper
43
44    return decorate
45
46
47# Example use
48@logged(logging.DEBUG)
49def add(x, y):
50    return x + y
51
52@logged(logging.CRITICAL, 'example')
53def spam():
54    print("Spam!")

运行示例

 1import logging
 2>>> logging.basicConfig(level=logging.DEBUG)
 3>>> add(2,3)
 4DEBUG:__main__:add
 55
 6>>> spam()
 7CRITICAL:example:spam
 8Spam!
 9>>> add.set_level(logging.WARNING)
10>>> add(5,6)
11WARNING:__main__:add
1211
13>>> add.set_message("add called")
14>>> add(4,5)
15WARNING:__main__:add called
169

本节的关键在于访问器函数(set_message()set_level()),它们以属性的方式附加到了包装函数上,每个访问器函数允许对nonlocal变量赋值来调整内部参数。

之所以要使用访问器函数,不使用对函数属性的直接访问,如下:

1@wraps(func)
2def wrapper(*args,**kwargs):
3    warpper.log.log(wrapper.level,wrapper.logmsg)
4    return func(*args,*kwargs)
5#Attach adjustment attributes
6wrapper.level = level
7wrapper.logmsg = logmsg
8wrapper.log = log

是因为这种方法只能用在最顶层的装饰器,如果当前顶层的装饰器上有添加了一个装饰器,这样就会隐藏下层的属性使得无法被修改,而使用访问器函数可以绕过这个限制。 本方法可以作为类装饰器的一种替代方案。