rate_limiter.py

class HeaderRateLimiter:
    """
    A Rate Limiter context manager that reads headers commonly used by APIs
    (such as Github, Twitter, ...).
    Can be used in Multi-Threaded environments.

    TODO: Add redis persistence to support multi-process/multi-machine

    Usage:
    rl = HeaderRateLimiter()
    for i in range(1000):
        with rl:
            resp = requests.head(url)
            rl.update(resp)
    """

    def __init__(self, limit_header="X-RateLimit-Limit", remaining_header="X-RateLimit-Remaining", reset_header="X-RateLimit-Reset"):
        self.lock = threading.Lock()
        self.limit_header = limit_header
        self.remaining_header = remaining_header
        self.reset_header = reset_header
        self.limit = 0
        self.remaining = 0
        self.reset = int(time.time())

    def __enter__(self):
        with self.lock:
            logging.debug(
                f"Limit: {self.limit} -- Remaining: {self.remaining} -- Reset: {self.reset}")
            if not self.remaining:
                reset_in = self.reset - int(time.time())
                if reset_in > 0:
                    logging.info(
                        f"Reached rate limit, waiting for {reset_in}s")
                    time.sleep(reset_in)
            self.remaining -= 1
        return self

    def __exit__(self, type, value, traceback):
        pass

    def update(self, resp):
        with self.lock:
            self.limit = int(resp.headers[self.limit_header])
            self.reset = int(resp.headers[self.reset_header])
            self.remaining = int(resp.headers[self.remaining_header])