Python 装饰器

Python 装饰器

装饰器的定义

装饰器本质上是一个Python函数,它可以让其他函数在不需要做任何代码变动的前提下增加额外功能,装饰器的返回值也是一个函数对象。 它经常用于有切面需求的场景,比如:插入日志、性能测试、事务处理、缓存、权限校验等场景。装饰器是解决这类问题的绝佳设计,有了装饰器,我们就可以抽离出大量与函数功能本身无关的雷同代码并继续重用。

为什么需要装饰器

假设你原先的程序实现了简单的 fun() 函数

#!/usr/bin/env python3
def fun():
    print('hello')
fun()

随后需要统计函数的运行时间,你当然可以对原函数进行修改,但是如果其他函数也实现该功能,逐个修改会添加大量重复代码。为了复用代码可以定义一个函数用于处理统计运行时间的需求。

#!/usr/bin/env python3
import time
from functools import wraps

def timethis(func):
    '''
    Decorator that reports the execution time.
    '''
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, end-start)
        return result
    return wrapper

@timethis
def fun():
    """
    hello
    """
    print('hello')
# 等同于
fun = timethis(fun)

fun()

通过上述函数对原函数进行包装,原函数 fun() 就向被装饰了一样,而timethis() 返回的是装饰后函数。

函数元数据

装饰器作用在某个函数时,这个函数重要的元数据,如名字、文档字符串、注解等信息都丢失了

# 未装饰前
>>> fun()
hello
>>> fun.__name__
'fun'
>>> fun.__doc__
'\n\thello\n\t'
>>> fun.__annotations__
{}
# 未使用 wraps 函数
>>> fun.__name__
'wrapper'
>>> fun.__doc__
>>> fun.__annotations__
{}
# 使用 wraps 函数
>>> fun.__name__
'fun'
>>> fun.__doc__
'\n\thello\n\t'
>>> fun.__annotations__
{}

可以看出函数的元数据被替换为 wrapper 函数对象的元数据。通过调用 @wraps() 装饰器,可以保留原函数的元数据

Note @wraps(func) 注解是很重要的,它能保留原始函数的元数据。可通过属性访问 __wrapped__ 被包装的函数。当多个装饰器作用于同一个函数时,__wrapped__ 属性调用的行为是不可预知的,应当避免这样做。

Note 并不是所有的装饰器都使用了 @wraps(),内置装饰器 @staticmethod@classmethod 是将原函数属性储存在 __func__ 中。

定义一个带参数的装饰器

假设你需要一个装饰器,给函数添加日志功能,同时可以指定日志的级别或其他选项。你可以这样定义装饰器

#!/usr/bin/env python3

from functools import wraps
import logging

def logged(level, name=None, message=None):
    """
    Add logging to a function.

    :param: level: logging level
    :param: name: log name
    :param: message: log message
    :return: wrapped func
    """
    def decorator(func):
        logname = name if name else func.__module__
        log = logging.getLogger(logname)
        logmsg = message if message else func.__name__

        @wraps(func)
        def wrapper(*args, **kwargs):
            log.log(level, logmsg)
            return func(*args, **kwargs)
        return wrapper
    return decorator

@logged(logging.DEBUG)
def fun():
    """
    print hello
    """
    print('hello')
# 等价于 func = logged(logging.DEBUG)(func)

@logged(logging.CRITICAL, 'example')
def fun():
    """
    print hello
    """
    print('hello')

decorate 返回一个可调用对象(包含 __call__ 方法),并传入被包装函数作为参数。

自定义属性的装饰器

使用一个装饰器包装一个函数,并允许运行时修改用户提供的参数。

#!/usr/bin/env python3

import logging
from functools import wraps, partial

def attach_wrapper(obj, func=None):
    """
    Utils decorator to attach a function as an attribute of obj
    """
    if func is None:
        return partial(attach_wrapper, obj)
    setattr(obj, func.__name__, func)
    return  func

def logged(level, name=None, message=None):
    """
    Add logging to a function.

    :param: level: logging level
    :param: name: log name
    :param: message: log message
    :return: wrapped func
    """
    def decorator(func):
        logname = name if name else func.__name__
        log = logging.getLogger(logname)
        logmsg = message if message else func.__name__

        @wraps(func)
        def wrapper(*args, **kwargs):
            log.log(level, logmsg)
            return func(*args, **kwargs)

        @attach_wrapper(wrapper)
        def set_level(newlevel):
            nonlocal level
            level = newlevel

        @attach_wrapper(wrapper)
        def set_message(newmsg)
            nonlocal logmsg
            logmsg = newmsg

        @attach_wrapper(wrapper)
        def get_level():
            return level

        return wrapper
    return decorator
@wraps(func)
def wrapper(*args, **kwargs):
    wrapper.log.log(wrapper.level, wrapper.logmsg)
    return func(*args, **kwargs)

# Attach adjustable attributes
wrapper.level = level
wrapper.logmsg = logmsg
wrapper.log = log

访问函数允许使用 nonlocal 来修改函数内部的变量。

直接修改的方法也可能正常工作,但前提是它必须是最外层的装饰器才行。 如果它的上面还有另外的装饰器,那么它会隐藏底层属性,使得修改它们没有任何作用。

可选参数装饰器

写一个装饰器,可以不传参数 @logged,或者是传递可选参数如 @logged(logging.DEBUG)

#!/usr/bin/env python3

import logging
from functools import wraps, partial

def logged(func=None, *, level=logging.DEBUG, name=None, message=None):
    """
    Add logging to a function.

    :param: func: function name
    :param: level: logging level
    :param: log message
    :return: wrapped func
    """
    if func is None:
        return partial(logged, level=level, name=name, message=message)

    logname = name if name else func.__name__
    log = logging.getLogger(logname)
    logmsg = message if message else func.__name__

    @wraps(func)
    def wrapper(*args, **kwargs):
        log.log(level, logmsg)
        return func(*args, **kwargs)

    return wraper

理解上述代码原理,需要知道装饰的调用规则,

# 简单装饰器

@logged
def fun():
    print('hello')
# 等价于 fun = logged(fun)

# 带参数装饰器

@logged(level=logging.DEBUG)
def fun():
    print('hello')
# 等价于 fun = logged(level=loggin.DEBUG)(fun)

简单装饰器将被包装函数作为第一个参数传递给 logged 函数,因此,logged() 中第一个参数就是被包装函数自身。

带参数装饰器调用时没有将被包装函数传递给 logged 函数,因此在装饰器内,它必须是可选的。这个反过来会迫使其他参数必须使用关键字来指定。 并且,但这些参数被传递进来后,装饰器要返回一个接受一个函数参数并包装它的函数。 为了这样做,我们使用了一个技巧,就是利用 functools.partial 。 它会返回一个未完全初始化的自身,除了被包装函数外其他参数都已经确定下来了。

利用装饰器进行参数检查

// TODO

类中定义装饰器

from functools import wraps

class A(object):
    """
    Decorator as an instance method
    """
    def decorator1(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            print("decorator1")
            reurn func(*args, **kwargs)
        return wrapper

    @classmethod
    def decorator2(cls, func):
        @wraps(func)
        def wraper(*args, **kwargs):
            print("decorator2")
            return func(*args, **kwargs)
        return wrapper

a = A()
@a.decorator1
def fun():
    print('hello')

@A.decorator2
def fun():
    print('hello')

可以看出一个是实例调用,另一个是类调用。装饰器要被定义成类方法并且你必须显式的使用父类名去调用它。

在类中定义装饰器初看上去好像很奇怪,但是在标准库中有很多这样的例子。 特别的,@property 装饰器实际上是一个类,它里面定义了三个方法 getter(), setter(), deleter(), 每一个方法都是一个装饰器。

它为什么要这么定义的主要原因是各种不同的装饰器方法会在关联的 property 实例上操作它的状态。 因此,任何时候只要你碰到需要在装饰器中记录或绑定信息,那么这不失为一种可行方法。

类作为装饰器

使用装饰器包装函数,但是希望返回一个可调用实例。为了实现类装饰器,你需要确保它实现了 __call__()__get__() 方法

import types
from functools import wraps

class Profiled(object):
    def __init__(self, func):
        wraps(func)(self)
        self.ncalls = 0

    def __call__(self, *args, **kwargs):
        self.ncalls += 1
        return self.__wrapped__(*args, **kwargs)

    def __get__(self, instance, cls):
        if instance is None:
            return self
        else
            return types.MethodType(self, instance)

@Profiled
def add(x, y):
    return x + y

class Spam:
    @Profiled
    def bar(self, x):
        print(self, x)

首先 wraps 函数将被包装函数的元信息复制到可调用实例中。

// TODO

类装饰器进阶

typys Method 方法