Define keybuilder protocol

This lets others create key builders that are type checked fully.

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
 

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]]]:
     """

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]:
+        ...