Python - 装饰器

装饰器是一种软件设计模式，可以动态的修改函数、方法、类的功能，而不需要修改原函数或者重写方法，PEP-0318 提出了对装饰器语法的支持，并在 Python 2.4 实现

嵌套函数

Python 中函数是可以嵌套的，我们可以写出以下代码

def parent_func(number):
    
    def chilren_func1():
        return 'hello, world, chilren1'
    
    def chilren_func2():
        return 'hello, world, chilren2'
    
    if number > 2:
        return chilren_func2
    return chilren_func1

如果我们直接传入数字并调用 parent_func，它会返回一个函数

>>> parent_func(3)
<function parent_func.<locals>.chilren_func2 at 0x104f0b048>

>>> parent_func(1)
<function parent_func.<locals>.chilren_func1 at 0x104ef3f28>

对其返回的函数再次调用

>>> parent_func(1)()
'hello, world, chilren1'

>>> parent_func(3)()
'hello, world, chilren2'

将函数作为参数

符合以下两个条件之一就算是高阶函数

将函数作为参数
将函数作为返回值

Python 中自然也可以支持将函数作为参数传入一个函数中，最经典的莫过于 map、reduce、sort、max、filter 等高阶函数

d 是一个 key 和 value 都是 int 的字典, 通过 max 函数取其最大值，默认只能通过对比 key 的值

d = {
  1: 2,
  2: 1,
  3: 0
}

>>> max(d)
3

但是可以传入一个函数，通过这个函数来进行对比, 这样我们就可以对比字典中 value 的值

>>> max(d, key=lambda x: d[x])
1

装饰器

有了以上基础后，就可以介绍真正的装饰器了

所谓 Python 的装饰器就是通过封装来动态修改一个函数或方法

我们可以传入一个函数到另一个函数中，另一个函数对其的功能进行修改或扩展，并将修改过的函数返回

def decroator_func(fn, *args, **kwargs):
    
    print('start {}'.format(fn.__name__))
    result = fn(*args, **kwargs)
    print('end {}'.format(fn.__name__))
    return result

def p_hello():
    print('hello, world')

>>> decroator_func(p_hello)
start p_hello
hello, world
end p_hello

但是上面这种方法不是正确的写法，我们需要嵌套一个 wrapper 函数来修改原函数，并返回修改后的函数，还需要通过一个变量来保存原函数的输出，由封装函数 reutrn

def decroator_func(fn):
    
    def wrapper(*args, **kwargs):
        print('start {}'.format(fn.__name__))
        result = fn(*args, **kwargs)
        print('end {}'.format(fn.__name__))
        return result
    return wrapper

def p_hello():
    print('hello, world')

p_hello = decroator_func(p_hello)

>>> p_hello()
start p_hello
hello, world
end p_hello

语法糖

虽然上面这种 p_hello = decroator_func(p_hello) 在功能上没有问题，但是看起来非常的丑陋，Python 可以在定义函数时通过 @ 这个语法糖来简化调用装饰器函数的步骤，使得代码更加简洁优雅

@decroator_func
def p_hello():
    print('hello, world')

functools.wraps

实际上，上面的代码还是有一些问题的，因为装饰器函数其实是返回 wrapper 函数，也就是封装过后的原函数，但是原函数的 docstring、__name__ 等属性都被 wrapper 覆盖了

>>> p_hello.__name__
'wrapper'

那么这时候我们就可以通过 functools 中的 wraps 这个装饰器函数来解决这个问题，我们的代码可以这样写

from functools import wraps

def decroator_func(fn):
    
    @wraps(fn)
    def wrapper(*args, **kwargs):
        print('start {}'.format(fn.__name__))
        result = fn(*args, **kwargs)
        print('end {}'.format(fn.__name__))
        return result
    return wrapper

再次尝试

@decroator_func
def p_hello():
    print('hello, world')

>>> p_hello.__name__
'p_hello'

带参数的装饰器

函数能接收参数，装饰器函数如果需要传入参数得多封装一层

from functools import wraps
from time import sleep

def sleep_decroator(second):

    def decroator(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            print('sleep {} s'.format(second))
            sleep(second)
            result = fn(*args, **kwargs)
            return result
        return wrapper
    return decroator

@sleep_decroator(3)
def p_hello():
    print('hello, world')
    
>>> p_hello()
sleep 3 s
hello, world
>>> p_hello.__name__
'p_hello'

类装饰器

Python 也可以使用类作为装饰器，可以让复杂的装饰器更加的优雅

其实装饰器只需要是一个 callable 对象即可，所以我们需要定义一个 __call__ 方法

class Decorator:
    def __init__(self):
        pass

    def __call__(self, fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            print('start {}'.format(fn.__name__))
            result = fn(*args, **kwargs)
            print('end {}'.format(fn.__name__))
            return result
        return wrapper

@Decorator()
def p_hello():
    print('hello, world')
    
>>> p_hello()
start p_hello
hello, world
end p_hello

上面这种情况下的代码看起来不是很好，但是如果我们定义一个需要传递参数的装饰器，代码就会很简洁

class SleepDecorator:
    def __init__(self, second):
        self.second = second

    def __call__(self, fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            print('sleep {} s'.format(self.second))
            sleep(self.second)
            result = fn(*args, **kwargs)
			print('end {}'.format(fn.__name__))
            return result
        return wrapper

@SleepDecorator(3)
def p_hello():
    print('hello, world')

>>> p_hello()
sleep 3 s
hello, world