Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: allow Requests to be sent to exempt_when #160

Merged
merged 13 commits into from
Jun 27, 2024
4 changes: 1 addition & 3 deletions slowapi/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -482,7 +482,7 @@ def __evaluate_limits(
limit_for_header = None
for lim in limits:
limit_scope = lim.scope or endpoint
if lim.is_exempt:
if lim.is_exempt(request):
continue
if lim.methods is not None and request.method.lower() not in lim.methods:
continue
Expand Down Expand Up @@ -699,11 +699,9 @@ def decorator(func: Callable[..., Response]):
else:
self._route_limits.setdefault(name, []).extend(static_limits)

connection_type: Optional[str] = None
sig = inspect.signature(func)
for idx, parameter in enumerate(sig.parameters.values()):
if parameter.name == "request" or parameter.name == "websocket":
connection_type = parameter.name
break
else:
raise Exception(
Expand Down
15 changes: 12 additions & 3 deletions slowapi/wrappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from typing import Callable, Iterator, List, Optional, Union

from limits import RateLimitItem, parse_many # type: ignore
from starlette.requests import Request


class Limit(object):
Expand All @@ -28,16 +29,24 @@ def __init__(
self.methods = methods
self.error_message = error_message
self.exempt_when = exempt_when
self._exempt_when_takes_request = len(inspect.signature(self.exempt_when).parameters) == 1
self.cost = cost
laurentS marked this conversation as resolved.
Show resolved Hide resolved
self.override_defaults = override_defaults

@property
def is_exempt(self) -> bool:
def is_exempt(self, request: Optional[Request] = None) -> bool:
"""
Check if the limit is exempt.

** parameter **
* **request**: the request object

Return True to exempt the route from the limit.
"""
return self.exempt_when() if self.exempt_when is not None else False
if self.exempt_when is None:
return False
if self._exempt_when_takes_request and request:
return self.exempt_when(request)
return self.exempt_when()

@property
def scope(self) -> str:
Expand Down
Loading