python/slowapi
1
0
Fork 0
mirror of https://github.com/laurentS/slowapi.git synced 2026-03-13 09:10:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

Enable dynamic limits dependant on key.

Browse Source
This commit is contained in:
Marat Sarbasov
2022-01-15 16:02:55 +03:00
parent d94c5eefad
commit 33899c0ef1
2 changed files with 19 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
slowapi/extension.py
View File

@@ -503,7 +503,7 @@ class Limiter:
if name in self._dynamic_route_limits:
for lim in self._dynamic_route_limits[name]:
try:
dynamic_limits.extend(list(lim))
dynamic_limits.extend(list(lim.with_request(request)))
except ValueError as e:
self.logger.error(
"failed to load ratelimit for view function %s (%s)",

23
slowapi/wrappers.py
View File

@@ -1,3 +1,4 @@
import inspect
from typing import Callable, Iterator, List, Optional, Union
from limits import RateLimitItem, parse_many # type: ignore
@@ -74,13 +75,21 @@ class LimitGroup(object):
self.error_message = error_message
self.exempt_when = exempt_when
self.override_defaults = override_defaults
self.request = None
def __iter__(self) -> Iterator[Limit]:
limit_items: List[RateLimitItem] = parse_many(
self.__limit_provider()
if callable(self.__limit_provider)
else self.__limit_provider
)
if callable(self.__limit_provider):
if "key" in inspect.signature(self.__limit_provider).parameters.keys():
assert (
"request" in inspect.signature(self.key_function).parameters.keys()
)
assert self.request
limit_raw = self.__limit_provider(self.key_function(self.request))
else:
limit_raw = self.__limit_provider()
else:
limit_raw = self.__limit_provider
limit_items: List[RateLimitItem] = parse_many(limit_raw)
for limit in limit_items:
yield Limit(
limit,
@@ -92,3 +101,7 @@ class LimitGroup(object):
self.exempt_when,
self.override_defaults,
)
def with_request(self, request):
self.request = request
return self