# Standard library imports
import contextlib
import math
import threading
import time
import sys
# Local imports
from uplink import decorators, utils
from uplink.clients.io import RequestTemplate, transitions
__all__ = ["ratelimit", "RateLimitExceeded"]
# Use monotonic time if available, otherwise fall back to the system clock.
now = time.monotonic if hasattr(time, "monotonic") else time.time
def _get_host_and_port(base_url):
parsed_url = utils.urlparse.urlparse(base_url)
return parsed_url.hostname, parsed_url.port
class RateLimitExceeded(RuntimeError):
"""A request failed because it exceeded the client-side rate limit."""
def __init__(self, calls, period):
super(RateLimitExceeded, self).__init__(
"Exceeded rate limit of [%s] calls every [%s] seconds."
% (calls, period)
)
class Limiter(object):
_last_reset = _num_calls = None
def __init__(self, max_calls, period, clock):
self._max_calls = max_calls
self._period = period
self._clock = clock
self._lock = threading.RLock()
self._reset()
@property
def period_remaining(self):
return self._period - (self._clock() - self._last_reset)
@contextlib.contextmanager
def check(self):
with self._lock:
if self.period_remaining <= 0:
self._reset()
yield self._max_calls > self._num_calls
self._num_calls += 1
def _reset(self):
self._num_calls = 0
self._last_reset = self._clock()
class RateLimiterTemplate(RequestTemplate):
def __init__(self, limiter, create_limit_reached_exception):
self._limiter = limiter
self._create_limit_reached_exception = create_limit_reached_exception
def before_request(self, request):
with self._limiter.check() as ok:
if ok:
return # Fallback to default behavior
elif self._create_limit_reached_exception is not None:
raise self._create_limit_reached_exception()
else:
return transitions.sleep(self._limiter.period_remaining)
# noinspection PyPep8Naming
class ratelimit(decorators.MethodAnnotation):
"""
A decorator that constrains a consumer method or an entire
consumer to making a specified maximum number of requests within a
defined time period (e.g., 15 calls every 15 minutes).
Note:
The rate limit is enforced separately for each host-port
combination. Logically, requests are grouped by host and port,
and the number of requests within a time period are counted and
capped separately for each group.
By default, when the limit is reached, the client will wait until
the current period is over before executing any subsequent
requests. If you'd prefer the client to raise an exception when the
limit is exceeded, set the ``raise_on_limit`` argument.
Args:
calls (int): The maximum number of allowed calls that the
consumer can make within the time period.
period (float): The duration of each time period in seconds.
raise_on_limit (:class:`Exception` or bool, optional): Either an
exception to raise when the client exceeds the rate limit or
a :class:`bool`. If :obj:`True`, a
:class:`~uplink.ratelimit.RateLimitExceeded` exception is
raised.
"""
BY_HOST_AND_PORT = _get_host_and_port
def __init__(
self,
calls=15,
period=900,
raise_on_limit=False,
group_by=BY_HOST_AND_PORT,
clock=now,
):
self._max_calls = max(1, min(sys.maxsize, math.floor(calls)))
self._period = period
self._clock = clock
self._limiter_cache = {}
self._group_by = utils.no_op if group_by is None else group_by
if utils.is_subclass(raise_on_limit, Exception) or isinstance(
raise_on_limit, Exception
):
self._create_limit_reached_exception = raise_on_limit
elif raise_on_limit:
self._create_limit_reached_exception = (
self._create_rate_limit_exceeded
)
else:
self._create_limit_reached_exception = None
def _get_limiter_for_request(self, request_builder):
key = self._group_by(request_builder.base_url)
try:
return self._limiter_cache[key]
except KeyError:
return self._limiter_cache.setdefault(
key, Limiter(self._max_calls, self._period, self._clock)
)
def modify_request(self, request_builder):
limiter = self._get_limiter_for_request(request_builder)
request_builder.add_request_template(
RateLimiterTemplate(limiter, self._create_limit_reached_exception)
)
def _create_rate_limit_exceeded(self):
return RateLimitExceeded(self._max_calls, self._period)
