Add request, response if it's not in the signautre

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ultra-cache"
-version = "0.1.0"
+version = "0.1.1"
 description = "Simple and extensible caching for FastAPI requests"
 authors = ["Jegor Kitskerkin <[email protected]>"]
 license = "MIT"
@@ -18,6 +18,7 @@ ruff = "^0.4.2"
 pyright = "^1.1.361"
 freezegun = "^1.5.0"
 pytest-mock = "^3.14.0"
+httpx = "^0.27.0"
 
 [build-system]
 requires = ["poetry-core"]

tests/test_decorator.py

@@ -9,10 +9,11 @@
 class FnWithArgs:
     "Utility class for testing and mocking"
 
-    def __init__(self, fn, args, kwargs) -> None:
+    def __init__(self, fn, args, kwargs, injected_kwargs) -> None:
         self.fn = fn
         self.args = args
         self.kwargs = kwargs
+        self.injected_kwargs = injected_kwargs
 
     def __str__(self) -> str:
         return f"{self.fn}(*{self.args}, **{self.kwargs})"
@@ -26,27 +27,19 @@ async def _fn_async(p1, p2):
     return p1 == p2
 
 
-def _fn_sync_request_1(r: Request, p1, p2):
+def _fn_sync_request_1(p1, p2, request: Request):
     return p1 == p2
 
 
-def _fn_sync_request_2(p1, p2, request: Request):
+def _fn_sync_request_2(p1, p2, r: Request):
     return p1 == p2
 
 
-def _fn_sync_request_3(p1, p2, r: Request):
+def _fn_sync_response_1(p1, p2, response: Response):
     return p1 == p2
 
 
-def _fn_sync_response_1(r: Response, p1, p2):
-    return p1 == p2
-
-
-def _fn_sync_response_2(p1, p2, response: Response):
-    return p1 == p2
-
-
-def _fn_sync_response_3(p1, p2, r: Response):
+def _fn_sync_response_2(p1, p2, r: Response):
     return p1 == p2
 
 
@@ -61,20 +54,56 @@ def key_builder():
 
 
 sample_args = (1, 2)
-sample_request = Request({"type": "http", "headers": {}})
-sample_response = Response()
+
+
+def sample_request():
+    return Request({"type": "http", "headers": {}})
+
+
+def sample_response():
+    return Response()
 
 
 @pytest.fixture(
     params=[
-        FnWithArgs(_fn_sync, sample_args, {}),
-        FnWithArgs(_fn_async, sample_args, {}),
-        FnWithArgs(_fn_sync_request_1, (sample_request, *sample_args), {}),
-        FnWithArgs(_fn_sync_request_2, sample_args, {"request": sample_request}),
-        FnWithArgs(_fn_sync_request_3, sample_args, {"r": sample_request}),
-        FnWithArgs(_fn_sync_response_1, (sample_response, *sample_args), {}),
-        FnWithArgs(_fn_sync_response_2, sample_args, {"response": sample_response}),
-        FnWithArgs(_fn_sync_response_3, sample_args, {"r": sample_response}),
+        FnWithArgs(
+            _fn_sync,
+            sample_args,
+            {},
+            {"request": sample_request(), "response": sample_response()},
+        ),
+        FnWithArgs(
+            _fn_async,
+            sample_args,
+            {},
+            {"request": sample_request(), "response": sample_response()},
+        ),
+        FnWithArgs(
+            _fn_sync_request_1,
+            sample_args,
+            {
+                "request": sample_request(),
+            },
+            {"response": sample_response()},
+        ),
+        FnWithArgs(
+            _fn_sync_request_2,
+            sample_args,
+            {"r": sample_request()},
+            {"response": sample_response()},
+        ),
+        FnWithArgs(
+            _fn_sync_response_1,
+            sample_args,
+            {"response": sample_response()},
+            {"request": sample_request()},
+        ),
+        FnWithArgs(
+            _fn_sync_response_2,
+            sample_args,
+            {"r": sample_response()},
+            {"request": sample_request()},
+        ),
     ]
 )
 def fn_with_args(request):
@@ -98,14 +127,20 @@ async def test_decorator_cached(
     )  # Weird syntax, but did not find any alternative
 
     cached_fn = cache(storage=storage)(fn_with_args.fn)
-    result = await cached_fn(*fn_with_args.args, **fn_with_args.kwargs)
+    result = await cached_fn(
+        *fn_with_args.args, **{**fn_with_args.kwargs, **fn_with_args.injected_kwargs}
+    )
 
-    # Note: no request and response in args/kwargs
-    key = key_builder(fn_with_args.fn, sample_args, kwargs={})
-    await cached_fn(*fn_with_args.args, **fn_with_args.kwargs)
+    # call second time
+    await cached_fn(
+        *fn_with_args.args, **{**fn_with_args.kwargs, **fn_with_args.injected_kwargs}
+    )
 
     assert spy_on_get.call_count == 2
 
+    # Note: no request and response in args/kwargs
+    key = key_builder(fn_with_args.fn, sample_args, kwargs={})
+
     spy_on_save.assert_called_once_with(key=key, value=(False), ttl=None)
     spy_on_fn.assert_called_once_with(*fn_with_args.args, **fn_with_args.kwargs)
 

tests/test_e2e.py

@@ -0,0 +1,23 @@
+from fastapi.testclient import TestClient
+from .utils import app
+
+client = TestClient(app)
+
+
+def test_cache_decorator():
+    response = client.get("/items/1")
+    assert response.status_code == 200
+    assert response.json() == {"item_id": 1}
+    assert response.headers.get("X-Cache") == "MISS"
+
+    # Test cache hit
+    response = client.get("/items/1")
+    assert response.status_code == 200
+    assert response.json() == {"item_id": 1}
+    assert response.headers.get("X-Cache") == "HIT"
+
+    # Test cache miss
+    response = client.get("/items/2")
+    assert response.status_code == 200
+    assert response.json() == {"item_id": 2}
+    assert response.headers.get("X-Cache") == "MISS"

tests/utils.py

@@ -0,0 +1,12 @@
+from fastapi import FastAPI
+from ultra_cache.decorator import cache
+from ultra_cache.storage.inmemory import InMemoryStorage
+
+app = FastAPI()
+storage = InMemoryStorage()
+
+
+@app.get("/items/{item_id}")
+@cache(storage=storage)
+async def read_item(item_id: int):
+    return {"item_id": item_id}

ultra_cache/decorator.py

@@ -15,16 +15,20 @@
 P = ParamSpec("P")
 R = TypeVar("R")
 
+S1 = TypeVar("S1")
+S2 = TypeVar("S2")
 
-def _extract_special_param(_fn, param_type: type) -> None | inspect.Parameter:
-    sig = inspect.signature(_fn)
+
+def _extract_param_of_type(
+    sig: inspect.Signature, param_type: type
+) -> None | inspect.Parameter:
     for _, param in sig.parameters.items():
         if param.annotation == param_type:
             return param
     return None
 
 
-def _extract[S1, S2](
+def _extract(
     param: inspect.Parameter | None, args: tuple[S1, ...], kwargs: dict[str, S2]
 ) -> tuple[tuple[S1, ...], dict[str, S2]]:
     if param is None:
@@ -39,15 +43,11 @@ def _extract[S1, S2](
             (i for i, p in enumerate(args) if isinstance(p, param.annotation)), -1
         )
         if request_index == -1:
-            raise ValueError(f"No argument of type {param.annotation} found in args")
+            return args, kwargs
         args_copy = args[:request_index] + args[request_index + 1 :]
         return (args_copy, kwargs)
 
 
-def _merge_args[T](value: T, index: int, *tup: T) -> tuple[T, ...]:
-    return tup[:index] + (value,) + tup[index:]
-
-
 def cache(
     ttl: int | float | None = None,
     build_cache_key: BuildCacheKey = DefaultBuildCacheKey(),
@@ -56,26 +56,35 @@ def cache(
     def _wrapper(
         func: Callable[P, Union[R, Coroutine[R, Any, Any]]],
     ) -> Callable[P, Coroutine[R, Any, Any]]:
-        # allows for the decorator to be used with fastapi params interospection
-        @wraps(func)
-        async def _decorator(*args: P.args, **kwargs: P.kwargs):
-            nonlocal storage
+        sig = inspect.signature(func)
 
-            request_param = _extract_special_param(func, Request)
-            response_param = _extract_special_param(func, Response)
+        original_request_param = _extract_param_of_type(sig, Request)
+        original_response_param = _extract_param_of_type(sig, Response)
+        request_param = original_request_param
+        response_param = original_response_param
 
-            response: Response | None = None
-            request: Request | None = None
+        new_parameters = list(sig.parameters.values())
+        if request_param is None:
+            request_param = inspect.Parameter(
+                "request", annotation=Request, kind=inspect.Parameter.KEYWORD_ONLY
+            )
+            new_parameters.append(request_param)
+        if response_param is None:
+            response_param = inspect.Parameter(
+                "response", annotation=Response, kind=inspect.Parameter.KEYWORD_ONLY
+            )
+            new_parameters.append(response_param)
 
-            if request_param is not None:
-                request = kwargs.get(request_param.name)
+        func.__signature__ = sig.replace(parameters=new_parameters)
 
-            if response_param is not None:
-                response = kwargs.get(response_param.name)
+        # allows for the decorator to be used with fastapi params interospection
+        @wraps(func)
+        async def _decorator(*args: P.args, **kwargs: P.kwargs):
+            nonlocal storage
+            request: Request = kwargs.get(request_param.name)
+            response: Response = kwargs.get(response_param.name)
 
-            cache_control = None
-            if request is not None:
-                cache_control = request.headers.get("Cache-Control", None)
+            cache_control = request.headers.get("Cache-Control", None)
 
             no_cache = False
             no_store = False
@@ -97,12 +106,15 @@ async def _decorator(*args: P.args, **kwargs: P.kwargs):
                 cached = await storage.get(key)
 
             if cached is not None:
-                if response:
-                    response.headers["X-Cache"] = "HIT"
+                response.headers["X-Cache"] = "HIT"
                 return cached
 
-            elif response:
-                response.headers["X-Cache"] = "MISS"
+            response.headers["X-Cache"] = "MISS"
+
+            if original_request_param is None:
+                kwargs.pop("request")
+            if original_response_param is None:
+                kwargs.pop("response")
 
             # Note: inspect.iscoroutinefunction returns False for AsyncMock
             if asyncio.iscoroutinefunction(func):