Python开源项目解读—ratelimit，限制函数单位时间内被调用次数

用户8949263

发布于 2023-11-09 17:17:28

5100

发布于 2023-11-09 17:17:28

文章被收录于专栏：Python数据分析实例

这个项目的开发背景是考虑一些服务的API 对于开放人员的访问频率会做一些限制，如果不小心超出了这个限制，服务可能会进制开发人员访问。

ratelimit 提供的装饰器，可以控制被装饰的函数在某个周期内被调用的次数不超过一个阈值，尽管作者本意是限制那些访问web API 的函数的调用次数，但你可以推而广之，所有不能频繁调用的函数都可以用这个装饰器来修饰。

项目的github地址: https://github.com/tomasbasham/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

被limits 装饰以后，call_api这个函数在15分钟内最多只能调用15次，超出后就会报错。

1. RateLimitDecorator

1.1 用类实现装饰器

我看了一下源码，作者的实现非常的简单，从ratelimit引入的limits 其实是一个类

limits = RateLimitDecorator
rate_limited = RateLimitDecorator # For backwards compatibility

来看一下RateLimitDecorator 这个类的实现

class RateLimitDecorator(object):
    '''
    Rate limit decorator class.
    '''
    def __init__(self, calls=15, period=900, clock=now(), raise_on_limit=True):
        self.clamped_calls = max(1, min(sys.maxsize, floor(calls)))
        self.period = period
        self.clock = clock
        self.raise_on_limit = raise_on_limit

        # Initialise the decorator state.
        self.last_reset = clock()
        self.num_calls = 0

        # Add thread safety.
        self.lock = threading.RLock()

    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kargs):
            with self.lock:
                period_remaining = self.__period_remaining()

                # If the time window has elapsed then reset.
                if period_remaining <= 0:
                    self.num_calls = 0
                    self.last_reset = self.clock()

                # Increase the number of attempts to call the function.
                self.num_calls += 1

                # If the number of attempts to call the function exceeds the
                # maximum then raise an exception.
                if self.num_calls > self.clamped_calls:
                    if self.raise_on_limit:
                        raise RateLimitException('too many calls', period_remaining)
                    return

            return func(*args, **kargs)
        return wrapper

    def __period_remaining(self):
        elapsed = self.clock() - self.last_reset
        return self.period - elapsed

这便是ratelimit 最核心部分的代码了，作者使用类实现了一个python装饰器，这种实现方法的关键是实现类的__call__方法。在进行装饰的时候，写法是这样的

@limits(calls=15, period=FIFTEEN_MINUTES)
def call_api(url):
    pass

这段代码等价于下面的写法

limits_decorator = RateLimitDecorator(calls=15, period=FIFTEEN_MINUTES)

call_api = limits_decorator(call_api)

limits_decorator 是RateLimitDecorator 类的一个实例，但由于RateLimitDecorator 实现了__call__方法，所以类的实例也是callable 的，因此limits_decorator(call_api) 等价于limits_decorator.call(call_api)，这便是使用类实现装饰器的原理所在。

1.2 线程锁

作者考虑到了多线程的场景，因此在wrapper函数加了线程锁，如果没有线程锁，多个线程同时修改self.num_calls 的值就可能导致调用次数记录的不准确。

RLock是可重入锁，关于线程锁，推荐你阅读我的教程python多线程

1.3 RateLimitException

作者自定义了一个异常类 RateLimitException，我们在工程实践时也应该多写一些自定义异常，这有助于我们在抛出异常时针对性的做处理。尽量不要在捕获异常时很笼统的捕获所有异常，那样虽然写起来简单，但不能促使我们更进一步的思考程序可能存在的问题。

自定义异常的方法很简单

class RateLimitException(Exception):
    '''
    Rate limit exception class.
    '''
    def __init__(self, message, period_remaining):
        super(RateLimitException, self).__init__(message)
        self.period_remaining = period_remaining

你可以定义新的初始化参数，记得调用super函数来进行初始化。

1.4 限制被调用次数的逻辑

装饰器在装饰函数时记录下当前的时间，这个动作对应在__init__函数中的self.last_reset = clock() 语句，当函数被调用时，self.__period_remaining() 会返回当前时间与self.last_reset的差值，如果小于零，说明还在周期时间内，如果此时调用次数超过了限制次数，就抛出异常。如果差值大于零，说明已经是一个新的限制周期了，重置self.last_reset 和 self.num_calls

3. 重试装饰器

def sleep_and_retry(func):
    @wraps(func)
    def wrapper(*args, **kargs):
        while True:
            try:
                return func(*args, **kargs)
            except RateLimitException as exception:
                time.sleep(exception.period_remaining)
    return wrapper

作者提供了sleep_and_retry装饰器与RateLimitDecorator一同使用，当RateLimitDecorator装饰的函数调用次数超出限制时会抛出异常RateLimitException，而RateLimitException 初始化时的第二个参数是这个周期内剩余的时间，在sleep_and_retry装饰器里，会根据这个时间sleep一段时间等待再次调用。

两个装饰器配合起来使用的方式

from ratelimit import limits, sleep_and_retry

import requests

FIFTEEN_MINUTES = 900

@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

先使用limits 对call_api 进行装饰，再用sleep_and_retry 进行二次装饰，一旦超出访问限制，程序不会结束，sleep_and_retry会根据当前访问周期剩余时间进行sleep ，然后再次调用。