带参数的Python装饰器让你的代码更优雅

引言

在上一篇文章中，我们介绍了 Python 装饰器的基本概念及其简单用法。

前面讲到的装饰器都是不带参数的装饰器，在需要对装饰器做一些针对性的处理的时候就不太适用了，这个时候需要对装饰器传入一些参数，根据传入的参数进行不同的处理。

带参数装饰器在实际开发中能够灵活地调整函数行为，广泛应用于日志记录、权限验证和缓存等场景。。

如何定义带参数的装饰器

带参数的装饰器的结构稍微复杂一些。我们要定义一个外层函数来接收装饰器的参数，并在这个外层函数中定义实际的装饰器函数，相当于将之前的装饰器用一个正常的函数包装了一层。下面来看一下带参数的装饰器和不带参数的装饰器的对比。

不带参数的装饰器

def decorator(func):
    def wrapper():
        print(f"Decorator")
        return func()

    return wrapper

带参数的装饰器：

def decorator_with_args(arg1=None, arg2=None):
    def decorator(func):
        def wrapper():
            print(f"Decorator arguments: {arg1}, {arg2}")
            return func()

        return wrapper

    return decorator

在上面带参数的装饰器中，decorator_with_args 是外层装饰器函数，它接收两个参数 arg1 和 arg2。内部的 decorator 函数是实际的装饰器，而 wrapper 函数则是用来包装被装饰的函数。

如何使用带参数的装饰器

@decorator_with_args("Hello", "World")
def greet():
    print(f"Greeting name")

greet()

输出：

Decorator arguments: Hello, World
Greeting name

在这个例子中，greet 函数被 decorator_with_args 装饰，传入了两个参数。装饰器在调用 greet 函数之前，先打印了装饰器的参数。

注意：带参数的装饰器有两个可选参数，可以选择不传。但需要注意的是，@decorator_with_args() 不能简写为 @decorator_with_args，因为 decorator_with_args 本质上是一个函数，必须调用后才能生效。

两种装饰器用法区别

使用不带参数的装饰器时，用法如下，要注意区别。

def decorator(func):
    def wrapper():
        print(f"Decorator")
        return func()

    return wrapper


@decorator
def greet():
    print(f"Greeting name")


greet()

处理可变参数

在很多时候，我们会遇到被装饰的函数有参数的情况，这时要怎么处理呢。

为了解决这个问题，我们可以使用 *args 和 **kwargs 来接收函数的所有参数。

装饰器处理可变参数

def func_decorator(decorator_arg):
    def decorator(func):
        def wrapper(*args, **kwargs):
            print(f"Decorator argument: {decorator_arg}")
            print(f"Function arguments: {args}, {kwargs}")
            return func(*args, **kwargs)

        return wrapper

    return decorator


@func_decorator("Test")
def add(*args):
    return sum(args)


result = add(1, 2, 3, 4)
print(f"Result: {result}")

输出：

Decorator argument: Test
Function arguments: (1, 2, 3, 4), {}
Result: 10

在这个示例中，func_decorator 装饰器可以处理在被装饰的函数有参数的情况，并且可以自适应任何参数，装饰器能够正确处理这些参数。

带参数装饰器的使用案例

自定义日志格式

在实际开发中，我们常常需要记录函数的执行日志。通过带参数的装饰器，我们可以自定义被装饰函数的参数和返回值，甚至可以做到修改被装饰函数的参数和返回值。

def log_decorator(fun_args=False, fun_result=False):
    def decorator(func):
        def wrapper(*args, **kwargs):
            if fun_args:
                print(f'func args is: {args} {kwargs}')
            if fun_result:
                result = func(*args, **kwargs)
                print(f'func result is: {result}')
                return result
            else:
                return func(*args, **kwargs)

        return wrapper

    return decorator


@log_decorator(True, True)
def multiply(x, y):
    return x * y


print(f"Result: {multiply(3, 5)}")


@log_decorator(False, False)
def multiply(x, y):
    return x * y


print(f"Result: {multiply(3, 5)}")

结果：

func args is: (3, 5) {}
func result is: 15
Result: 15
---
Result: 15

在这个示例中，log_decorator 装饰器接受两个布尔型参数，用于判断是否要输出函数参数和函数返回值。

此处日志输出比较简单，仅做为示例参考，生产环境要做一些比较复杂的判断并且要将日志持久化到磁盘的日志文件中。

缓存机制

import time
from functools import lru_cache


class Timer:
    def __init__(self):
        self.start = None

    def __enter__(self):
        import time
        self.start = time.time()

    def __exit__(self, type, value, traceback):
        print(f'time is:{time.time() - self.start}')


@lru_cache(maxsize=20)
def fibonacci(n):
    if n <= 1:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)


with Timer() as timer:
    print(fibonacci(15))

with Timer() as timer2:
    def fibonacci(n):
        if n <= 1:
            return n
        return fibonacci(n - 1) + fibonacci(n - 2)


    print(fibonacci(15))

结果：

14 13
13 12
12 11
11 10
10 9
9 8
8 7
7 6
6 5
5 4
4 3
3 2
2 1
1 0
610
time is:4.076957702636719e-05
--------------------
14 13
13 12
12 11
11 10
10 9
9 8
8 7
7 6
--- 此处数据量太大，省略，完整的输出大约 1000 行，大家可以自行运行查看结果
1 0
610
time is:0.002106189727783203

此处我使用了 functools 包中提供的 lru_cache 缓存装饰器，用于缓存在计算斐波那契数列时重复的计算值。

从运行结果可以看到，不使用装饰器时，函数进行了大量的重复计算，导致最后的运行耗时和使用了装饰器的结果差了好几个数量级。

总结

带参数的装饰器为我们提供了极大的灵活性，使得我们可以根据不同的需求来调整函数的行为。在实际应用中，带参数的装饰器可以用于多种场景，如自定义日志、缓存、控制函数的执行行为等等。