Python 实现函数调用次数限制：基于装饰器的限流方案

核心思路：用装饰器实现限流

在 Python 开发中，限制函数调用频率是一个常见需求，例如防止 API 被恶意刷取或控制资源消耗。本文的方案灵感来源于常见的“函数调用次数缓存”装饰器，但将核心功能从“缓存结果”替换为“限制访问次数”。通过闭包维护状态，我们可以在不修改原函数逻辑的前提下，为其添加限流保护。

代码实现：限流装饰器

以下代码定义了一个名为 stat_called_time 的装饰器。它通过字典 cache 记录每个被装饰函数的调用次数和最后更新时间。当函数被调用时，装饰器会检查当前时间与上次记录的时间差是否小于设定的时间窗口（此处示例为 60 秒）。如果在时间窗口内，调用次数累加；否则重置计数。当调用次数超过预设的限制值（示例中为 10 次）时，函数将不再执行原逻辑，而是返回 None 并打印提示信息。

#newtest.py
#!/usr/bin/env python
#-*- coding:utf-8 -*-
import time

def stat_called_time(func):
    cache = {}
    limit_times = [10]  # 使用列表以便在闭包中修改整数值
    
    def _called_time(*args, **kwargs):
        key = func.__name__
        
        if key in cache.keys():
            call_times, updatetime = cache[key]
            # 判断是否在时间窗口内（60秒）
            if time.time() - updatetime < 60:
                cache[key][0] += 1
            else:
                # 超出时间窗口，重置计数和时间
                cache[key] = [1, time.time()]
        else:
            # 首次调用，初始化计数和时间
            call_times = 1
            cache[key] = [call_times, time.time()]
            
        print('调用次数: %s' % cache[key][0])
        print('限制次数: %s' % limit_times[0])
        
        # 判断是否超过限制次数
        if cache[key][0] <= limit_times[0]:
            res = func(*args, **kwargs)
            cache[key][1] = time.time()  # 更新最后调用时间
            return res
        else:
            print("超过调用次数了")
            return None
            
    return _called_time

@stat_called_time
def foo():
    print("I'm foo")

if __name__ == '__main__':
    for i in range(10):
        foo()

测试验证：跨模块调用与时间窗口测试

为了验证限流效果，我们创建另一个脚本 test.py 导入上述模块并进行测试。测试分为几个阶段：首先连续调用 30 次，观察限流触发情况；接着暂停并等待一段时间，验证时间窗口重置机制；最后再次调用，确认计数是否已重置。

#test.py
from newtest import foo
import time

# 第一阶段：快速调用 30 次
for i in range(30):
    foo()
print('*'*20)
foo()
foo()
print('*'*20)

# 第二阶段：等待 60 秒，使时间窗口过期
for i in range(60):
    print(i)
    time.sleep(1)

# 第三阶段：时间窗口重置后，再次调用 11 次
for i in range(11):
    foo()

关键点解析

• 闭包状态维护： 装饰器内部使用字典 cache 存储每个函数的调用状态。由于 Python 中整数是不可变对象，这里使用列表 limit_times = [10] 来存储限制值，以便在闭包中通过索引修改其内容。
• 时间窗口判断： 通过 time.time() - updatetime < 60 判断当前时间是否仍在限流窗口内。若在窗口内，则累加调用次数；若超出窗口，则重置计数器和时间戳。
• 限流执行： 当调用次数超过限制值时，直接返回 None 并打印警告信息，不再执行原函数逻辑。

以上就是 Python 限制函数调用次数实例的详细内容，更多关于 Python 装饰器限流与函数调用的资料请关注本站其它相关文章！