diff --git a/fastapi_limiter/__init__.py b/fastapi_limiter/__init__.py index 2b644a2..5b48091 100644 --- a/fastapi_limiter/__init__.py +++ b/fastapi_limiter/__init__.py @@ -1,5 +1,5 @@ from math import ceil -from typing import Callable, Union +from typing import Callable, Optional, Union from fastapi import HTTPException from starlette.requests import Request @@ -46,11 +46,11 @@ async def ws_default_callback(ws: WebSocket, pexpire: int): class FastAPILimiter: redis = None - prefix: str = None - lua_sha: str = None - identifier: Callable = None - http_callback: Callable = None - ws_callback: Callable = None + prefix: Optional[str] = None + lua_sha: Optional[str] = None + identifier: Optional[Callable] = None + http_callback: Optional[Callable] = None + ws_callback: Optional[Callable] = None lua_script = """local key = KEYS[1] local limit = tonumber(ARGV[1]) local expire_time = ARGV[2] @@ -76,7 +76,7 @@ async def init( identifier: Callable = default_identifier, http_callback: Callable = http_default_callback, ws_callback: Callable = ws_default_callback, - ): + ) -> None: cls.redis = redis cls.prefix = prefix cls.identifier = identifier @@ -85,5 +85,5 @@ async def init( cls.lua_sha = await redis.script_load(cls.lua_script) @classmethod - async def close(cls): + async def close(cls) -> None: await cls.redis.close() diff --git a/fastapi_limiter/depends.py b/fastapi_limiter/depends.py index 58de9f5..43fa0eb 100644 --- a/fastapi_limiter/depends.py +++ b/fastapi_limiter/depends.py @@ -1,6 +1,6 @@ -from typing import Callable, Optional +from typing import Annotated, Callable, Optional -from pydantic import conint +from pydantic import Field from starlette.requests import Request from starlette.responses import Response from starlette.websockets import WebSocket @@ -12,11 +12,11 @@ class RateLimiter: def __init__( self, - times: conint(ge=0) = 1, - milliseconds: conint(ge=-1) = 0, - seconds: conint(ge=-1) = 0, - minutes: conint(ge=-1) = 0, - hours: conint(ge=-1) = 0, + times: Annotated[int, Field(ge=0)] = 1, + milliseconds: Annotated[int, Field(ge=-1)] = 0, + seconds: Annotated[int, Field(ge=-1)] = 0, + minutes: Annotated[int, Field(ge=-1)] = 0, + hours: Annotated[int, Field(ge=-1)] = 0, identifier: Optional[Callable] = None, callback: Optional[Callable] = None, ):