diff --git a/fastapi_cache/__init__.py b/fastapi_cache/__init__.py index 42c585a..2393916 100644 --- a/fastapi_cache/__init__.py +++ b/fastapi_cache/__init__.py @@ -1,8 +1,9 @@ -from typing import Callable, ClassVar, Optional, Type +from typing import ClassVar, Optional, Type from fastapi_cache.backends import Backend from fastapi_cache.coder import Coder, JsonCoder from fastapi_cache.key_builder import default_key_builder +from fastapi_cache.types import KeyBuilder class FastAPICache: @@ -11,7 +12,7 @@ class FastAPICache: _expire: ClassVar[Optional[int]] = None _init: ClassVar[bool] = False _coder: ClassVar[Optional[Type[Coder]]] = None - _key_builder: ClassVar[Optional[Callable]] = None + _key_builder: ClassVar[Optional[KeyBuilder]] = None _enable: ClassVar[bool] = True @classmethod @@ -21,7 +22,7 @@ class FastAPICache: prefix: str = "", expire: Optional[int] = None, coder: Type[Coder] = JsonCoder, - key_builder: Callable = default_key_builder, + key_builder: KeyBuilder = default_key_builder, enable: bool = True, ) -> None: if cls._init: @@ -64,7 +65,7 @@ class FastAPICache: return cls._coder @classmethod - def get_key_builder(cls) -> Callable: + def get_key_builder(cls) -> KeyBuilder: assert cls._key_builder, "You must call init first!" # nosec: B101 return cls._key_builder diff --git a/fastapi_cache/decorator.py b/fastapi_cache/decorator.py index b0601aa..091da3f 100644 --- a/fastapi_cache/decorator.py +++ b/fastapi_cache/decorator.py @@ -2,7 +2,7 @@ import inspect import logging import sys from functools import wraps -from typing import Any, Awaitable, Callable, Optional, Type, TypeVar +from typing import Awaitable, Callable, Optional, Type, TypeVar if sys.version_info >= (3, 10): from typing import ParamSpec @@ -15,6 +15,7 @@ from starlette.responses import Response from fastapi_cache import FastAPICache from fastapi_cache.coder import Coder +from fastapi_cache.types import KeyBuilder logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -25,7 +26,7 @@ R = TypeVar("R") def cache( expire: Optional[int] = None, coder: Optional[Type[Coder]] = None, - key_builder: Optional[Callable[..., Any]] = None, + key_builder: Optional[KeyBuilder] = None, namespace: Optional[str] = "", ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]: """ diff --git a/fastapi_cache/types.py b/fastapi_cache/types.py new file mode 100644 index 0000000..a05a37e --- /dev/null +++ b/fastapi_cache/types.py @@ -0,0 +1,21 @@ +from typing import Any, Awaitable, Callable, Optional, Protocol, Union + +from starlette.requests import Request +from starlette.responses import Response + + +_Func = Callable[..., Any] + + +class KeyBuilder(Protocol): + def __call__( + self, + _function: _Func, + _namespace: str = ..., + *, + request: Optional[Request] = ..., + response: Optional[Response] = ..., + args: Optional[tuple[Any, ...]] = ..., + kwargs: Optional[dict[str, Any]] = ..., + ) -> Union[Awaitable[str], str]: + ...