Decorator进阶指南

在[python装饰器完全指南基础篇中]，我们已经知道了python中的装饰器本质上只是一个接受一个函数对象作为输入，在该函数前后增加一些新逻辑，并返回包装过后的新函数对象的函数。

问题一： 有输入参数的函数的装饰器

在基础篇中，我们所讨论的所有被装饰器装饰的函数都是不接收输入参数的函数。那么对于有输入参数的函数，我们应该如何定义适用于它们的装饰器呢？请看如下代码片段。

def decorator_for_func_with_arguments(func_to_decorate):
    def wrapper_that_passes_through_arguments(arg1, arg2):
        print('I got args! Look: {} {}'.format(arg1, arg2))
        func_to_decorate(arg1, arg2)
    return wrapper_that_passes_through_arguments

@ decorator_for_func_with_arguments
def print_full_name(first_name, last_name):
    print('I am {} {}'.format(first_name, last_name))

print_full_name('Tony', 'Stark')
# output: 
# I got args! Look: Tony Stark
# I am Tony Stark

可以看出，由于装饰器事实上返回了一个新的函数来代替我们需要装饰的函数，所以我们只需要确保在装饰器中返回的新的函数与原函数所接受的参数格式一致即可。

问题二：类方法的装饰器

Python中的类方法事实上与函数一样，只是固定接收当前实例的引用作为第一个参数，一般标记为self。那么能够装饰类方法的装饰器事实上也可以用与问题一中一致的方法实现，只不过我们总是要确保返回的函数所接收的第一个参数也是当前实例的引用（self）即可。如下所示。

def decorator_for_instance_method(method_to_decorate):
    def wrapper(self, bonus):
        # 升职加薪，奖金增加一倍 d=====(￣▽￣*)b
        bonus = bonus * 2
        return method_to_decorate(self, bonus)
    return wrapper

class Salary(object):
    def __init__(self):
        self.base = 666

    @decorator_for_instance_method
    def total_compensation(self, bonus):
        print('Congrats! You got a total compensation of {}'.format(self.base * 12 + bonus))

salary_instance = Salary()
salary_instance.total_compensation(2048)
# output: Congrats! You got a total compensation of 12088

类似地，我们也可以用python中的args和*kwargs来实现一个能够装饰接收任意数目参数函数的装饰器。如下所示。

def decorator_passing_arbitrary_arguments(function_to_decorate):
    def wrapper_with_arbitrary_arguments(*args, **kwargs):
        print('Received arguments as following')
        print(args)
        print(kwargs)

        function_to_decorate(*args, **kwargs)

    return wrapper_with_arbitrary_arguments

@decorator_passing_arbitrary_arguments
def function_with_no_argument():
    print('This function does not have any argument')

function_with_no_argument()
# output:
# Received arguments as following
# ()
# {}
# This function does not have any argument

@decorator_passing_arbitrary_arguments
def function_with_arguments(a, b, c):
    print('This function has arguments')

function_with_arguments(1,2,3)
# output:
# Received arguments as following
# (1, 2, 3)
# {}
# This function has arguments

@decorator_passing_arbitrary_arguments
def function_with_named_arguments(a, b, c, name)
    print('{}, {}, {}'.format(a, b, c))
    print('{}'.format(name))

function_with_named_arguments(1, 2, 3, name='robot')
# output:
# Received arguments as following
# (1, 2, 3)
# {'name': 'robot'}
# 1, 2, 3
# robot

class Salary(object):
    def __init__(self):
        self.base = 666

    @decorator_passing_arbitrary_arguments
    def total_compensation(self, bonus):
        print('Congrats! You got a total compensation of {}'.format(self.base * 12 + bonus))

salary = Salary()
salary.total_compensation(2048)
# salary.total_compensation(2048)
# Received arguments as following
# (<__main__.Salary object at 0x1070b5f28>, 2048)
# {}
# Congrats! You got a total compensation of 10040

问题三：给装饰器传入参数

上面我们讨论了装饰器所装饰的函数参数有关的问题。接下来我们讨论如何实现能够接收参数的装饰器。

看过之前文章内容的读者相比已经知道，所谓装饰器其实就是接收一个函数作为输入，并返回另一个函数的函数。这种情况下，由于装饰器的函数签名已经固定，所以我们无法直接传入除输入函数之外的参数。如下面代码所示。

# 装饰器就是接收一个函数作为输入的函数
def my_decorator(func):
    print('This is an ordinary function')
    def wrapper():
        print('This is the wrapper function that will be returned')
        func()

# 只有上述格式签名的函数才能作为装饰器
# 下面代码等于
# lazy_function = my_decorator(lazy_function)
@my_decorator
def lazy_function():
    print('zzzzz')

# output:
# This is an ordinary function

lazy_function()
# output:
# This is the wrapper function that will be returned
# zzzzz

上面代码说明当我们使用@my_decorator的时候，事实上就是执行了lazy_function = my_decorator(lazy_function)。因此在无法直接改变装饰器签名的情况下，我们需要采用一些别的办法来实现我们的目标——实现一个能够返回装饰器的函数来替装饰器接收参数，并使用闭包（对闭包概念不熟悉的读者，请参阅这篇文章
）的方法来将这些参数传递到装饰器中。换句话说，我们需要一个装饰器工厂函数来替我们动态生成装饰器。

首先我们实现一个装饰器工厂函数，如下面的代码所示。

def decorator_maker():
    print('This is a factory generating decorators on the fly. This function is called once we want a fresh decorator.')

    def my_decorator(func):
        print('This is the decorator generated on the fly. This function is called when the decoration happens.')
        # 类似地，我们还是在装饰器中定义一个wrapper还包装原函数
        def wrapper():
            print('This is the wrapper around the decorated function. This function is called once the decorated function is called.')
            return func()

        return wrapper

    print('The decorator function created on the fly is returned.')
    return my_decorator

def func():
    print('This is the function decorated.')

fresh_decorator = decorator_maker()
# output:
# This is a factory generating decorators on the fly. This function is called once we want a fresh decorator.
# The decorator function created on the fly is returned.

func = fresh_decorator(func)
# output:
# This is the decorator generated on the fly. This function is called when the decoration happens.

func()
# output:
# This is the wrapper around the decorated function. This function is called once the decorated function is called.
# This is the function decorated.

基于如上代码，我们可以用更加pythonic的方式去使用这个装饰器工厂函数。如下所示。

@decorator_maker()
def func():
    print('This is the function decorated.')
# output:
# This is a factory generating decorators on the fly. This function is called once we want a fresh decorator.
# The decorator function created on the fly is returned.
# This is the decorator generated on the fly. This function is called when the decoration happens.

func()
# output:
# This is the wrapper around the decorated function. This function is called once the decorated function is called.
# This is the function decorated.

上面的例子说明，我们是事实上可以用一个装饰器工厂返回的函数作为@语法中的装饰器使用。既然如此，我们自然也就可以给这个工厂函数传入参数。如下面代码所示。

def decorator_maker_with_arguments(decorator_arg1, decorator_arg2):
    print('This is the decorator factory. Input arguments are: {}, {}.'.format(decorator_arg1, decorator_arg2))

    def my_decorator(func):
        # 利用闭包特性来将工厂函数接收的参数传递给动态生成的装饰器
        print('This is the decorator function created on the fly.')
        print('Arguments are passed in from outer function using closure: {}, {}.'.format(decorator_arg1, decorator_arg2))

        # 注意这里wrapper函数接收的参数是被该装饰器装饰的函数
        # 不要和上面工厂函数接收的参数混淆
        def wrapper(function_arg1, function_arg2):
            print('This is the wrapper around the decorated function.')
            print('This function can access all the variables.')
            print('Variables from the decorator factory: {}, {}.'.format(decorator_arg1, decorator_arg2))
            print('Variables from the decorated function: {}, {}.'.format(function_arg1, function_arg2))
            return func(function_arg1, function_arg2)

        return wrapper
    return my_decorator

@decorator_maker_with_arguments('Avengers', 'Justice League')
def func(function_arg1, function_arg2):
    print('This is the function decorated.')
    print('It accepts two arguments: {}, {}.'.format(function_arg1, function_arg2))

# output:
# This is the decorator factory. Input arguments are: Avengers, Justice League.
# This is the decorator function created on the fly.
# Arguments are passed in from outer function using closure: Avengers, Justice League.

func('Captain America', 'Bat Man')
# output:
# This is the wrapper around the decorated function.
# This function can access all the variables.
# Variables from the decorator factory: Avengers, Justice League.
# Variables from the decorated function: Captain America, Bat Man.
# This is the function decorated.
# It accepts two arguments: Captain America, Bat Man.

上面例子中，我们将字符串常量作为参数传递给了装饰器工厂函数。与平常的python函数一样，这一函数也可以接收变量作为参数。如下面所示。

a1 = 'Avenagers'
a2 = Justice League'
b1 = 'Captain America'
b2 = 'Bat Man'

@decorator_maker_with_arguments(a1, a2)
def func(function_arg1, function_arg2):
    print('This is the function decorated.')
    print('It accepts two arguments: {}, {}.'.format(function_arg1, function_arg2))

# output:
# This is the decorator factory. Input arguments are: Avengers, Justice League.
# This is the decorator function created on the fly.
# Arguments are passed in from outer function using closure: Avengers, Justice League.

func(b1, b2)
# output:
# This is the wrapper around the decorated function.
# This function can access all the variables.
# Variables from the decorator factory: Avengers, Justice League.
# Variables from the decorated function: Captain America, Bat Man.
# This is the function decorated.
# It accepts two arguments: Captain America, Bat Man.

通过上述讨论，可以看出对于装饰器工厂函数而言，调用它与调用普通函数完全相同。我们甚至也可以使用*args及**kwargs来接受可变长度的参数。但有一点需要注意，那就是装饰器装饰函数的过程只发生一次，那就是python编译器import当前代码文件的时候。因此一旦文件已经被import过，那我们就无法修改装饰器工厂接收的参数，也无法修改已经生成的装饰器装饰过的函数了。

Reference
文中部分内容翻译自如下文章。翻译部分版权归原作者所有。
https://gist.github.com/Zearin/2f40b7b9cfc51132851a