Python 装饰器

函数对象

在 Python 中，函数是一级对象，它可以赋值给变量，也可以作为参数传递给另外一个函数。

函数赋值给变量

函数是一个对象，可以被赋值给变量。

def hi():
    print("Hi, James.")

print(hi)    # <function hi at 0x7fe4b9001510>
hi()         # Hi, James.

hello = hi

print(hello) # <function hi at 0x7fe4b9001510>
hello()      # Hi, James.

函数作为参数

函数可以作为参数传递给另一个函数。

def hi():
    print('Hi, James.')
    
def before_hi(func):
    print('Before excuting hi()')
    func()

before_hi(hi)
# Before excuting hi()
# Hi, James.

函数作为返回值

函数不仅可以作为参数，也可以作为另一个函数的返回值。

def hi():
    return 'Hi'

def hello():
    return hi

h = hello()
h()  # 'Hi'

函数中定义函数

我们可以将函数定在另外一个函数中，也就是嵌套函数。

def hi():
    print('Hi')
    
    def hello():
        print('hello')
    
    hello()
    
hi()
# Hi
# hello

装饰器

理解了函数是一级对象后，装饰器也就呼之欲出了。

装饰器其实就是个可调用对象（函数或带有 __call__ 的类等），它接收一个函数，完成一些操作后，返回另外一个函数。

使用装饰器，我们可以在不修改原来代码的基础上，更直观的为代码添加额外的功能，如权限验证、日志记录等。

语法

我们编写一个简单的装饰器：

def decorated_by(func):
    return func

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

add = decorated_by(add)

add(2, 3)  # 5

从 Python 2.5 开始，在函数声明前使用 @ 应用装饰器：

def decorated_by(func):
    print('I am a decorator')
    return func

@decorated_by      # I am a decorator
def add(x, y):
    return x + y
    
add(2, 3)  # 5

@decorated_by 实际上干这么一件事：add = decorated_by(add)，也就是被装饰的函数依然存在，只是变量 add 指向了新的函数。

更重要的是可以看出：在应用装饰器到函数时，执行装饰器代码。

多个装饰器

使用多个装饰器时，按照从下到上的顺序应用它们

def decorated_a(func):
    print('I am decorator a')
    return func

def decorated_b(func):
    print('I am decorator b')
    return func

def decorated_c(func):
    print('I am decorator c')
    return func

@decorated_a
@decorated_b
@decorated_c
def add(x, y):
    return x + y

# I am decorator c
# I am decorator b
# I am decorator a

编写装饰器

我们来编写一个装饰器，确保函数接收到的参数都是整型。

def requires_ints(func):

    def wrapper(*args, **kwargs):
        values = [i for i in kwargs.values()]
        for arg in list(args) + values:
            if not isinstance(arg, int):
                raise TypeError('Only accept integers.')
        return func(*args, **kwargs)
    
    return wrapper


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

print(add.__name__)  # wrapper
print(add.__doc__)   # None

可以发现 wrapper 替换原函数的 name、docstring，这并不是我们想要的结果。使用 functools.wraps 来解决这个问题：

from functools import wraps

def requires_ints(func):
    
    @wraps(func)
    def wrapper(*args, **kwargs):
        values = [i for i in kwargs.values()]
        for arg in list(args) + values:
            if not isinstance(arg, int):
                raise TypeError('Only accept integers.')
        return func(*args, **kwargs)
    
    return wrapper

@requires_ints
def add(x, y):
    '''Return the sum of x and y'''
    return x + y

print(add.__name__)  # add
print(add.__doc__)   # Return the sum of x and y

在 Python 2 中使用 help(add) 查看时，显示的参数信息仍为 *args 和 **kwargs。

带参数的装饰器

上面写的装饰器看上去并没有任何参数（使用 @ 时被装饰的方法作为隐式参数传递给装饰器）。但是，有时候让装饰器自带一些参数，可以提供更灵活的应用。

我们改写 requires_ints，使其还可以限制被装饰函数的参数的数量：

from functools import wraps

def requires_ints(count=0):
    
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            values = [i for i in kwargs.values()]
            args_list = list(args) + [i for i in kwargs.values()]
            
            if count == 0:
                # Not limit
                pass
            elif len(args_list) > count:
                raise Exception(f'The number of arguments cannot more than {count}')
            
            for arg in args_list:
                if not isinstance(arg, int):
                    raise TypeError('Only accept integers.')
            return func(*args, **kwargs)

        return wrapper
    return decorator


@requires_ints(5)
def add(*args):
    '''Return the sum of all arguments'''
    return sum(args)

分解下上述装饰器的运行步骤：

1. @requires_ints() 返回函数 decorator
2. add(*args) 作为参数传递进 decorator
3. 之后就和前面的所写的装饰器一样了，带参数的装饰器只是多了一层嵌套。

我们之前应用装饰器时使用的是 @requires_ints 形式，而在这里需使用 @requires_ints() 返回真正的装饰器，然后才能发挥效果，这显然是不太友好的。

下面我们来改进一下，使 @requires_ints 和 @requires_ints() 都可以。

from functools import wraps

def requires_ints(_decorated=None, count=0):
    
    if _decorated and count:
        raise RuntimeError('Unexpected  arguments.')
    
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            values = [i for i in kwargs.values()]
            args_list = list(args) + [i for i in kwargs.values()]
            if count == 0:
                # Not limit
                pass
            elif len(args_list) > count:
                raise Exception(f'The number of arguments cannot more than {count}')
            
            for arg in args_list:
                if not isinstance(arg, int):
                    raise TypeError('Only accept integers.')
            return func(*args, **kwargs)

        return wrapper
    
    if _decorated:
        # 使用 @requires_ints 形式, 被装饰的方法赋值给 _decorated
        # 执行 decorator(_decorated) 返回 wrapper, 即和不带参数的装饰器一样
        return decorator(_decorated)
    else:
        # 使用 @requires_ints() 形式, _decorated 确保为 None, 不能由用户手动传入
        # 前面做了检测，如果用户手动传入了 _decorated 和 count, 则报错
        return decorator


@requires_ints
def add1(*args):
    '''Return the sum of all arguments'''
    return sum(args)

@requires_ints()
def add2(*args):
    '''Return the sum of all arguments'''
    return sum(args)

@requires_ints(count=5)
def add3(*args):
    '''Return the sum of all arguments'''
    return sum(args)