跳至主要內容

装饰器应用

Mr.Liu大约 4 分钟Python

熟悉装饰器的常见语法之后,可以利用装饰器解决平常遇到的问题。

functools.wraps

Python 装饰器(decorator)在实现的时候,被装饰后的函数其实已经是另外一个函数了(函数名等函数属性会发生改变),为了不影响原函数,Python 的 functools 包中提供了一个叫 wraps 的 decorator 来消除这样的副作用。写一个 decorator 的时候,最好在实现之前加上 functoolswraps,它能保留原有函数的名称和函数属性

不加 wraps

def my_decorator(func):
    def wrapper(*args, **kwargs):
        '''decorator'''
        print('Calling decorated function...')
        return func(*args, **kwargs)
    return wrapper

@my_decorator
def example():
    """Docstring"""
    print('Called example function')

print(example.__name__, example.__doc__)

运行结果

wrapper decorator

加上 wraps

import functools


def my_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        '''decorator'''
        print('Calling decorated function...')
        return func(*args, **kwargs)
    return wrapper

@my_decorator
def example():
    """Docstring"""
    print('Called example function')

print(example.__name__, example.__doc__)

运行结果

example Docstring

限频

请实现一个装饰器,限制该函数被调用的频率,如10秒1次

import functools
import time


def set_time(t, n):
    def set_num(func):
        dic = {"last_time": 0, "time_interval": t, "num": 0}

        @functools.wraps(func)
        def call_func():
            now_time = time.time()

            finall_time = dic["last_time"] + dic["time_interval"]  # 代表可以重新调用的时间
            finall_num = dic["num"]  # 代表调用的次数

            if finall_num < n:  # 当次数不满足时,可以继续调用
                if finall_num == 0:
                    dic["last_time"] = now_time
                dic["num"] += 1

                return func()

            elif now_time >= finall_time:  # 已经超出时间,可以重新调用
                dic["num"] = 0
                dic["last_time"] = 0

                return func()
            else:
                print("还有%.2fs才能调用该函数" % (finall_time - now_time))

        return call_func

    return set_num

下面测试一下:

import time


@set_time(10, 1)
def s():
    print("hello...")


s()
s()
time.sleep(2)
s()

结果:

hello...
还有10.00s才能调用该函数
还有8.00s才能调用该函数

超时退出

我们日常在使用的各种网络请求库时都带有timeout参数,例如request库。这个参数可以使请求超时就不再继续了,直接抛出超时错误,避免等太久。

如果我们自己开发的方法也希望增加这个功能,该如何做呢?

方法很多,但最简单直接的是使用并发库futures,为了使用方便,我将其封装成了一个装饰器,代码如下:

import functools
from concurrent import futures

executor = futures.ThreadPoolExecutor(1)

def timeout(seconds):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kw):
            future = executor.submit(func, *args, **kw)
            return future.result(timeout=seconds)
        return wrapper
    return decorator

定义了以上函数,我们就有了一个超时结束的装饰器,下面可以测试一下:

import time

@timeout(1)
def task(a, b):
    time.sleep(1.2)
    return a+b

task(2, 3)

结果:

Traceback (most recent call last):
  File "/Users/lsf/PycharmProjects/django/djangotest/book/demo.py", line 41, in <module>
    task(2, 3)
  File "/Users/lsf/PycharmProjects/django/djangotest/book/demo.py", line 30, in wrapper
    return future.result(timeout=self.seconds)
  File "/usr/local/Cellar/python@3.7/3.7.12/Frameworks/Python.framework/Versions/3.7/lib/python3.7/concurrent/futures/_base.py", line 437, in result
    raise TimeoutError()
concurrent.futures._base.TimeoutError

上面我们通过装饰器定义了函数的超时时间为1秒,通过睡眠模拟函数执行超过1秒时,成功的抛出了超时异常。

程序能够在超时时间内完成时:

@timeout(1)
def task(a, b):
    time.sleep(0.9)
    return a+b

task(2, 3)

结果:

5

可以看到,顺利的得到了结果。

这样我们就可以通过一个装饰器给任何函数增加超时时间,这个函数在规定时间内还处理不完就可以直接结束任务。

前面我将这个装饰器将所需的变量定义到了外部,其实我们还可以通过类装饰器进一步封装,代码如下:

import functools
from concurrent import futures

class timeout:
    __executor = futures.ThreadPoolExecutor(1)

    def __init__(self, seconds):
        self.seconds = seconds

    def __call__(self, func):
        @functools.wraps(func)
        def wrapper(*args, **kw):
            future = self.__executor.submit(func, *args, **kw)
            return future.result(timeout=self.seconds)
        return wrapper

经测试使用类装饰器能得到同样的效果。

日志记录

如果我们需要记录部分函数的执行时间,函数执行前后打印一些日志,装饰器是一种很方便的选择。

代码如下:

import time
import functools
 
def log(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        res = func(*args, **kwargs)
        end = time.perf_counter()
        print(f'函数 {func.__name__} 耗时 {(end - start) * 1000} ms')
        return res
    return wrapper

装饰器 log 记录某个函数的运行时间,并返回其执行结果。

测试一下:

@log
def now():
    print('2021-10-6')
    
now()

结果:

2021-10-6
函数 now 耗时 0.09933599994838005 ms

缓存

如果经常调用一个函数,而且参数经常会产生重复,如果把结果缓存起来,下次调用同样参数时就会节省处理时间。

定义函数:

import math
import random
import time


def task(x):
    time.sleep(0.01)
    return round(math.log(x**3 / 15), 4)

执行:

for i in range(500):
    task(random.randrange(5, 10))

结果:

Wall time: 5.91 s

此时如果我们使用缓存的效果就会大不一样,实现缓存的装饰器有很多,我就不重复造轮子了,这里使用 functools 包下的 LRU 缓存:

from functools import lru_cache

@lru_cache()
def task(x):
    time.sleep(0.01)
    return round(math.log(x**3 / 15), 4)

执行:

for i in range(500):
    task(random.randrange(5, 10))

结果:

Wall time: 0.05 ms