From f5015d539e319ffcd8a4b9a7e8891c67d0d754b3 Mon Sep 17 00:00:00 2001 From: Scott Date: Sat, 8 Jan 2022 15:21:25 -0800 Subject: [PATCH] Add rate limiter decorator. --- decorator_utils.py | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/decorator_utils.py b/decorator_utils.py index 4f98a6d..480543a 100644 --- a/decorator_utils.py +++ b/decorator_utils.py @@ -2,7 +2,6 @@ """Decorators.""" -import datetime import enum import functools import inspect @@ -15,7 +14,7 @@ import sys import threading import time import traceback -from typing import Callable, Optional +from typing import Any, Callable, Optional import warnings # This module is commonly used by others in here and should avoid @@ -58,6 +57,31 @@ def invocation_logged(func: Callable) -> Callable: return wrapper_invocation_logged +def rate_limited(n_per_second: int) -> Callable: + """Limit invocation of a wrapped function to n calls per second. + Thread safe. + + """ + min_interval = 1.0 / float(n_per_second) + + def wrapper_rate_limited(func: Callable) -> Callable: + last_invocation_time = [0.0] + + def wrapper_wrapper_rate_limited(*args, **kargs) -> Any: + while True: + elapsed = time.clock_gettime(0) - last_invocation_time[0] + wait_time = min_interval - elapsed + if wait_time > 0.0: + time.sleep(wait_time) + else: + break + ret = func(*args, **kargs) + last_invocation_time[0] = time.clock_gettime(0) + return ret + return wrapper_wrapper_rate_limited + return wrapper_rate_limited + + def debug_args(func: Callable) -> Callable: """Print the function signature and return value at each call.""" -- 2.46.0