wip

Author: HansKallekleiv

backend_py/primary/primary/middleware/add_browser_cache.py

@@ -1,10 +1,12 @@
 from functools import wraps
 from contextvars import ContextVar
-from typing import Dict, Any
+from typing import Dict, Any, Callable, TypeVar, Awaitable
 
 from starlette.datastructures import MutableHeaders
+from starlette.types import ASGIApp, Scope, Receive, Send, Message
 from primary.config import DEFAULT_CACHE_MAX_AGE, DEFAULT_STALE_WHILE_REVALIDATE
 
+T = TypeVar("T")
 
 # Initialize with a factory function to ensure a new dict for each context
 def get_default_context() -> Dict[str, Any]:
@@ -14,7 +16,9 @@ def get_default_context() -> Dict[str, Any]:
 cache_context: ContextVar[Dict[str, Any]] = ContextVar("cache_context", default=get_default_context())
 
 
-def add_custom_cache_time(max_age: int, stale_while_revalidate: int = 0) -> Any:
+def add_custom_cache_time(
+    max_age: int, stale_while_revalidate: int = 0
+) -> Callable[[Callable[..., Awaitable[T]]], Callable[..., Awaitable[T]]]:
     """
     Decorator that sets a custom cache time for the endpoint response.
 
@@ -28,9 +32,9 @@ async def my_endpoint():
             return {"data": "some_data"}
     """
 
-    def decorator(func):
+    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
         @wraps(func)
-        async def wrapper(*args, **kwargs):
+        async def wrapper(*args: Any, **kwargs: Any) -> T:
             context = cache_context.get()
             context["max_age"] = max_age
             context["stale_while_revalidate"] = stale_while_revalidate
@@ -47,17 +51,17 @@ class AddBrowserCacheMiddleware:
     Adds cache-control to the response headers
     """
 
-    def __init__(self, app):
+    def __init__(self, app: ASGIApp) -> None:
         self.app = app
 
-    async def __call__(self, scope, receive, send) -> None:
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http":
             return await self.app(scope, receive, send)
 
         # Set initial context and store token
         cache_context.set(get_default_context())
 
-        async def send_with_cache_header(message) -> None:
+        async def send_with_cache_header(message: Message) -> None:
             if message["type"] == "http.response.start":
                 headers = MutableHeaders(scope=message)
                 context = cache_context.get()