限制调用次数的 Python 装饰器

Question

我想编写一个装饰器来限制对包装函数的调用次数。比如说，如果我希望包装函数最多被调用 10 次，那么装饰器应该在前 10 次执行该函数，然后它应该返回 None。

这是我想出的：

from functools import wraps


def max_calls(num):
     """Decorator which allows its wrapped function to be called `num` times"""

     def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            calls = getattr(wrapper, 'calls', 0)
            calls += 1

            if calls == num:
                return None

            setattr(wrapper, 'calls', calls)
            return func(*args, **kwargs)
        setattr(wrapper, 'calls', 0)
        return wrapper
    return decorator

然而，这个计数正确调用，当达到限制时返回None，但是......它不会在程序运行之间重置。也就是说，如果我执行一次程序，计数器达到5，然后重新执行程序，它从5继续。我需要更改什么才能使装饰器正常工作？

Answer 1

问题在于您只维护一组呼叫计数。但这意味着每个 Flask 请求与所有其他请求共享调用计数。您需要做的是为每个 Flask 请求维护一组单独的调用计数。

从阅读API documentation 看来，完成这项工作的方法似乎是执行以下三个步骤：

1. 创建一个flask.Request 的子类，可以存储您的函数调用计数：

   import collections
   import flask
   
   class MyRequest(flask.Request):
       """A subclass of Request that maintains function call counts."""
       def __init__(self, *args, **kwargs):
           super(MyRequest, self).__init__(*args, **kwargs)
           self.call_counts = collections.defaultdict(int)
   

2. 在初始化应用程序时将request_class 设置为您的子类：

   app = flask.Flask(__name__, request_class=MyRequest, ...)
   

3. 重写您的装饰器以将其计数存储在全局 flask.request 对象中：

   import functools
   
   def max_calls(num, default=None):
       """Decorator which allows its wrapped function to be called at most `num`
       times per Flask request, but which then returns `default`.
   
       """
       def decorator(func):
           @functools.wraps(func)
           def wrapper(*args, **kwargs):
               if flask.request.call_counts[func] == num:
                   return default
               flask.request.call_counts[func] += 1
               return func(*args, **kwargs)
           return wrapper
       return decorator
   

但是写完之后，如果我不指出你的问题似乎很奇怪，那就太失职了。为什么要限制一个函数可以被 Flask 请求调用的次数？您是否尝试进行某种速率限制？使用其他方法似乎可以更好地完成您想做的任何事情。