Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.