学会这12个Python装饰器，让你的代码更上一层楼

学会这12个Python装饰器，让你的代码更上一层楼

Python 装饰器是个强大的工具，可帮你生成整洁、可重用和可维护的代码。某种意义上说，会不会用装饰器是区分新手和老鸟的重要标志。如果你不熟悉装饰器，你可以将它们视为将函数作为输入并在不改变其主要用途的情况下扩展其功能的函数。装饰器可以有效提高你的工作效率并避免重复代码。本文我整理了项目中经常用到的 12 个装饰器，值得每一个Python开发者掌握。

-

文章目录

• • 1. @logger
  • 2. @wraps
  • 3. @lru_cache
  • 4. @repeat
  • 5. @timeit
  • 6. @retry
  • 7. @countcall
  • 8. @rate_limited
  • 9. @dataclass
  • 10. @register
  • 11. @property
  • 12. @singledispatch
  • 结论

1. @logger

我们从最简单的装饰器开始，手动实现一个可以记录函数开始和结束的装饰器。被修饰函数的输出结果如下所示：

some_function(args)

# ----- some_function: start -----
# some_function executing
# ----- some_function: end -----

要实现一个装饰器，首先要给装饰器起一个合适的名称：这里我们给装饰器起名为logger。

装饰器本质上是一个函数，它将一个函数作为输入并返回一个函数作为输出。 输出函数通常是输入的扩展版。 在我们的例子中，我们希望输出函数用start和end语句包围输入函数的调用。

由于我们不知道输入函数都带有什么参数，我们可以使用 *args 和 **kwargs 从包装函数传递它们。*args 和 **kwargs 允许传递任意数量的位置参数和关键字参数。

下面是logger装饰器的示例代码：

def logger(function):
    def wrapper(*args, **kwargs):
        print(f"----- {function.__name__}: start -----")
        output = function(*args, **kwargs)
        print(f"----- {function.__name__}: end -----")
        return output
    return wrapper

logger函数可以应用于任意函数，比如：

decorated_function = logger(some_function)

上面的语句是正确的，但Python 提供了更 Pythonic 的语法——使用 @ 修饰符。因此更通常的写法是：

@logger
def some_function(text):
    print(text)

some_function("first test")
# ----- some_function: start -----
# first test
# ----- some_function: end -----

some_function("second test")
# ----- some_function: start -----
# second test
# ----- some_function: end -----

2. @wraps

此装饰器更新wrapper 函数，使其看起来像一个原始函数，并继承其名字和属性。

要了解 @wraps 的作用以及为什么需要它，让我们将前面写的logger装饰器应用到一个将两个数字相加的简单函数中。

下面的代码是未使用@wraps装饰器的版本：

def logger(function):
    def wrapper(*args, **kwargs):
        """wrapper documentation"""
        print(f"----- {function.__name__}: start -----")
        output = function(*args, **kwargs)
        print(f"----- {function.__name__}: end -----")
        return output
    return wrapper

@logger
def add_two_numbers(a, b):
    """this function adds two numbers"""
    return a + b

如果我们用__name__ 和 __doc__来查看被装饰函数add_two_numbers的名称和文档，会得到如下结果：

add_two_numbers.__name__
'wrapper'

add_two_numbers.__doc__
'wrapper documentation'

输出的是wrapper函数的名称和文档。这是我们预期想要的结果，我们希望保留原始函数的名称和文档。这时@wraps装饰器就派上用场了。

我们唯一需要做的就是给wrapper函数加上@wraps装饰器。

from functools import wraps

def logger(function):
    @wraps(function)
    def wrapper(*args, **kwargs):
        """wrapper documentation"""
        print(f"----- {function.__name__}: start -----")
        output = function(*args, **kwargs)
        print(f"----- {function.__name__}: end -----")
        return output
    return wrapper

@logger
def add_two_numbers(a, b):
    """this function adds two numbers"""
    return a + b

再此检查add_two_numbers函数的名称和文档，我们可以看到该函数的元数据。

add_two_numbers.__name__
# 'add_two_numbers'

add_two_numbers.__doc__
# 'this function adds two numbers'

3. @lru_cache

@lru_cache是Python内置装饰器，可以通过from functools import lru_cache引入。@lru_cache的作用是缓存函数的返回值，当缓存装满时，使用least-recently-used（LRU）算法丢弃最少使用的值。

@lru_cache装饰器适合用于输入输出不变且运行时间较长的任务，例如查询数据库、请求静态页面或一些繁重的处理。

在下面的示例中，我使用@lru_cache来修饰一个模拟某些处理的函数。然后连续多次对同一输入应用该函数。

import random
import time
from functools import lru_cache

​

@lru_cache(maxsize=None)
def heavy_processing(n):
    sleep_time = n + random.random()
    time.sleep(sleep_time)

# 初次调用
%%time
heavy_processing(0)
# CPU times: user 363 µs, sys: 727 µs, total: 1.09 ms
# Wall time: 694 ms

# 第二次调用
%%time
heavy_processing(0)
# CPU times: user 4 µs, sys: 0 ns, total: 4 µs
# Wall time: 8.11 µs

# 第三次调用
%%time
heavy_processing(0)
# CPU times: user 5 µs, sys: 1 µs, total: 6 µs
# Wall time: 7.15 µs

从上面的输出可以看到，第一次调用花费了694ms，因为执行了time.sleep()函数。后面两次调用由于参数相同，直接返回缓存值，因此并没有实际执行函数内容，因此非常快地得到函数返回。

4. @repeat

该装饰器的所用是多次调用被修饰函数。这对于调试、压力测试或自动化多个重复任务非常有用。

跟前面的装饰器不同，@repeat接受一个输入参数，

def repeat(number_of_times):
    def decorate(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for _ in range(number_of_times):
                func(*args, **kwargs)
        return wrapper
    return decorate

上面的代码定义了一个名为repeat的装饰器，有一个输入参数number_of_times。与前面的案例不同，这里需要decorate函数来传递被修饰函数。然后，装饰器定义一个名为wrapper的函数来扩展被修饰函数。

@repeat(5)
def hello_world():
    print("hello world")

hello_world()
# hello world
# hello world
# hello world
# hello world
# hello world

5. @timeit

该装饰器用来测量函数的执行时间并打印出来。这对调试和监控非常有用。

在下面的代码片段中，@timeit装饰器测量process_data函数的执行时间，并以秒为单位打印所用的时间。

import time
from functools import wraps

def timeit(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        end = time.perf_counter()
        print(f'{func.__name__} took {end - start:.6f} seconds to complete')
        return result
    return wrapper

@timeit
def process_data():
    time.sleep(1)

process_data()
# process_data took 1.000012 seconds to complete

6. @retry

当函数遇到异常时，该装饰器会强制函数重试多次。它接受三个参数：重试次数、捕获的异常以及重试之间的间隔时间。

其工作原理如下：

• wrapper函数启动num_retrys次迭代的for循环。
• 将被修饰函数放到try/except块中。每次迭代如果调用成功，则中断循环并返回结果。否则，休眠sleep_time秒后继续下一次迭代。
• 当for循环结束后函数调用依然不成功，则抛出异常。

示例代码如下：

import random
import time
from functools import wraps

def retry(num_retries, exception_to_check, sleep_time=0):
    """
    遇到异常尝试重新执行装饰器
    """
    def decorate(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for i in range(1, num_retries+1):
                try:
                    return func(*args, **kwargs)
                except exception_to_check as e:
                    print(f"{func.__name__} raised {e.__class__.__name__}. Retrying...")
                    if i < num_retries:
                        time.sleep(sleep_time)
            # 尝试多次后仍不成功则抛出异常
            raise e
        return wrapper
    return decorate

@retry(num_retries=3, exception_to_check=ValueError, sleep_time=1)
def random_value():
    value = random.randint(1, 5)
    if value == 3:
        raise ValueError("Value cannot be 3")
    return value

random_value()
# random_value raised ValueError. Retrying...
# 1

random_value()
# 5

7. @countcall

@countcall用于统计被修饰函数的调用次数。这里的调用次数会缓存在wraps的count属性中。

from functools import wraps

def countcall(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        wrapper.count += 1
        result = func(*args, **kwargs)
        print(f'{func.__name__} has been called {wrapper.count} times')
        return result
    wrapper.count = 0
    return wrapper

@countcall
def process_data():
    pass

process_data()
process_data has been called 1 times
process_data()
process_data has been called 2 times
process_data()
process_data has been called 3 times

8. @rate_limited

@rate_limited装饰器会在被修饰函数调用太频繁时，休眠一段时间，从而限制函数的调用速度。这在模拟、爬虫、接口调用防过载等场景下非常有用。

import time
from functools import wraps

def rate_limited(max_per_second):
    min_interval = 1.0 / float(max_per_second)
    def decorate(func):
        last_time_called = [0.0]
        @wraps(func)
        def rate_limited_function(*args, **kargs):
            elapsed = time.perf_counter() - last_time_called[0]
            left_to_wait = min_interval - elapsed
            if left_to_wait > 0:
                time.sleep(left_to_wait)
            ret = func(*args, **kargs)
            last_time_called[0] = time.perf_counter()
            return ret
        return rate_limited_function
    return decorate

该装饰器的工作原理是：测量自上次函数调用以来所经过的时间，并在必要时等待适当的时间，以确保不超过速率限制。其中等待时间=min_interval - elapsed，这里min_intervalue是两次函数调用之间的最小时间间隔（以秒为单位），已用时间是自上次调用以来所用的时间。如果经过的时间小于最小间隔，则函数在再次执行之前等待left_to_wait秒。

⚠注意：该函数在调用之间引入了少量的时间开销，但确保不超过速率限制。

如果不想自己手动实现，可以用第三方包，名叫ratelimit。

pip install ratelimit

使用非常简单，只需要装饰被调用函数即可:

from ratelimit import limits

import requests

FIFTEEN_MINUTES = 900

@limits(calls=15, period=FIFTEEN_MINUTES)
def call_api(url):
    response = requests.get(url)

    if response.status_code != 200:
        raise Exception('API response: {}'.format(response.status_code))
    return response

如果被装饰函数的调用次数超过允许次数，则会抛出ratelimit.RateLimitException异常。要处理该异常可以将@sleep_and_retry装饰器与@limits装饰器一起使用。

@sleep_and_retry
@limits(calls=15, period=FIFTEEN_MINUTES)
def call_api(url):
    response = requests.get(url)

    if response.status_code != 200:
        raise Exception('API response: {}'.format(response.status_code))
    return response

这样被装饰函数在再次执行之前会休眠剩余时间。

9. @dataclass

Python 3.7 引入了@dataclass装饰器，将其加入到标准库，用于装饰类。它主要用于存储数据的类自动生成诸如__init__， __repr__， __eq__， __lt__，__str__ 等特殊函数。这样可以减少模板代码，并使类更加可读和可维护。

另外，@dataclass还提供了现成的美化方法，可以清晰地表示对象，将其转换为JSON格式，等等。

from dataclasses import dataclass, 

@dataclass
class Person:
    first_name: str
    last_name: str
    age: int
    job: str

    def __eq__(self, other):
        if isinstance(other, Person):
            return self.age == other.age
        return NotImplemented

    def __lt__(self, other):
        if isinstance(other, Person):
            return self.age < other.age
        return NotImplemented

​

john = Person(first_name="John", 
              last_name="Doe", 
              age=30, 
              job="doctor",)

anne = Person(first_name="Anne", 
              last_name="Smith", 
              age=40, 
              job="software engineer",)

print(john == anne)
# False

print(anne > john)
# True

asdict(anne)
#{'first_name': 'Anne',
# 'last_name': 'Smith',
# 'age': 40,
# 'job': 'software engineer'}

10. @register

如果你的Python脚本意外终止，但你仍想执行一些任务来保存你的工作、执行清理或打印消息，那么@register在这种情况下非常方便。

from atexit import register

@register
def terminate():
    perform_some_cleanup()
    print("Goodbye!")

while True:
    print("Hello")

运行上面的代码会不断在控制台输出"Hello"，点击Ctrl + C强制终止脚本运行，你会看到控制台输出"Goodbye"，说明程序在中断后执行了@register装饰器装饰的terminate()函数。

11. @property

@property装饰器用于定义类属性，这些属性本质上是类实例属性的getter、setter和deleter方法。

通过使用@property装饰器，可以将方法定义为类属性，并将其作为类属性进行访问，而无需显式调用该方法。

如果您想在获取或设置值时添加一些约束和验证逻辑，使用@property装饰器会非常方便。

下面的示例中，我们在rating属性上定义了一个setter，对输入执行约束（介于0和5之间）。

class Movie:
    def __init__(self, r):
        self._rating = r

    @property
    def rating(self):
        return self._rating

    @rating.setter
    def rating(self, r):
        if 0 <= r <= 5:
            self._rating = r
        else:
            raise ValueError("The movie rating must be between 0 and 5!")

batman = Movie(2.5)
batman.rating
# 2.5

batman.rating = 4
batman.rating
# 4

batman.rating = 10

# ---------------------------------------------------------------------------
# ValueError                                Traceback (most recent call last)
# Input In [16], in ()
# ----> 1 batman.rating = 10
# Input In [11], in Movie.rating(self, r)
#      12     self._rating = r
#      13 else:
# ---> 14     raise ValueError("The movie rating must be between 0 and 5!")
#
# ValueError: The movie rating must be between 0 and 5!

12. @singledispatch

@singledispatch允许函数对不同类型的参数有不同的实现，有点像Java等面向对象语言中的函数重载。

from functools import singledispatch

@singledispatch
def fun(arg):
    print("Called with a single argument")

@fun.register(int)
def _(arg):
    print("Called with an integer")

@fun.register(list)
def _(arg):
    print("Called with a list")

fun(1)  # Prints "Called with an integer"
fun([1, 2, 3])  # Prints "Called with a list"

结论

装饰器是一个重要的抽象思想，可以在不改变原始代码的情况下扩展代码，如缓存、自动重试、速率限制、日志记录，或将类转换为超级数据容器等。

装饰器的功能远不止于此，本文介绍的12个常用装饰器只是抛砖引玉，当你理解了装饰器思想和用法后，可以发挥创造力，实现各种自定义装饰器来解决具体问题。

最后给大家推荐一个很棒的装饰器列表，里面记录了大量实用的、有趣的装饰器，大家可以多多尝试使用。