Skip to content

Commit

Permalink
Add type hints
Browse files Browse the repository at this point in the history
  • Loading branch information
lloesche committed Sep 5, 2024
1 parent a741e0b commit 0cb0601
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions fixbackend/auth/rate_limiter.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from datetime import timedelta
from textwrap import dedent
from typing import List, Optional, cast, Any
from fixbackend.types import Redis

from fixcloudutils.util import utc
Expand Down Expand Up @@ -28,14 +29,16 @@ def _new_tokens(self, tokens: int, ttl: int) -> float:
return min(self.limit, tokens + time_passed * self.refill_rate)

async def check(self, username: str) -> bool:
[tokens, ttl] = await self.redis.eval(
result = await self.redis.eval(
""" local ttl = redis.call('TTL', KEYS[1])
local tokens = redis.call('GET', KEYS[1])
return {tokens, ttl}
""",
1,
f"rate_limit:{username}",
)
result = cast(List[Optional[int]], result)
tokens, ttl = result

if tokens is None:
return True
Expand All @@ -45,7 +48,7 @@ async def check(self, username: str) -> bool:
return new_tokens >= 1

async def consume(self, username: str) -> bool:
allowed = await self.redis.eval(
allowed: Any = await self.redis.eval(
dedent(
"""
local bucket_key = KEYS[1]
Expand Down Expand Up @@ -87,5 +90,4 @@ async def consume(self, username: str) -> bool:
utc().timestamp(),
)

allowed = bool(allowed)
return allowed
return bool(allowed)

0 comments on commit 0cb0601

Please sign in to comment.