From c10806e0de86ed701cfa699dbe5cb21b3d694b7d Mon Sep 17 00:00:00 2001 From: Luna Date: Tue, 13 Aug 2019 12:24:03 -0300 Subject: [PATCH] run: add support for x-ratelimit-precision header --- run.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/run.py b/run.py index dbf1965..4700d1f 100644 --- a/run.py +++ b/run.py @@ -194,6 +194,21 @@ async def app_after_request(resp): return resp +def _set_rtl_reset(bucket, resp): + reset = bucket._window + bucket.second + precision = request.headers.get('x-ratelimit-precision', 'second') + + if precision == 'second': + resp.headers['X-RateLimit-Reset'] = str(round(reset)) + elif precision == 'millisecond': + resp.headers['X-RateLimit-Reset'] = str(reset) + else: + resp.headers['X-RateLimit-Reset'] = ( + 'Invalid X-RateLimit-Precision, ' + 'valid options are (second, millisecond)' + ) + + @app.after_request async def app_set_ratelimit_headers(resp): """Set the specific ratelimit headers.""" @@ -205,9 +220,8 @@ async def app_set_ratelimit_headers(resp): resp.headers['X-RateLimit-Limit'] = str(bucket.requests) resp.headers['X-RateLimit-Remaining'] = str(bucket._tokens) - resp.headers['X-RateLimit-Reset'] = str(bucket._window + bucket.second) - resp.headers['X-RateLimit-Global'] = str(request.bucket_global).lower() + _set_rtl_reset(bucket, resp) # only add Retry-After if we actually hit a ratelimit retry_after = request.retry_after