Add response extension (#4)

karpetrosyan · Jul 28, 2023 · 5450398 · 5450398
1 parent 3c77c29
commit 5450398
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/hishel/_async/_pool.py b/hishel/_async/_pool.py
@@ -42,6 +42,7 @@ async def handle_async_request(self, request: Request) -> Response:
 
             if isinstance(res, Response):
                 # Simply use the response if the controller determines it is ready for use.
+                res.extensions["from_cache"] = True  # type: ignore[index]
                 return res
 
             if isinstance(res, Request):
@@ -55,6 +56,7 @@ async def handle_async_request(self, request: Request) -> Response:
                     old_response=stored_resposne, new_response=response
                 )
                 await self._storage.store(key, response)
+                response.extensions["from_cache"] = response.status == 304  # type: ignore[index]
                 return response
 
         response = await self._pool.handle_async_request(request)
@@ -63,6 +65,7 @@ async def handle_async_request(self, request: Request) -> Response:
             await response.aread()
             await self._storage.store(key, response)
 
+        response.extensions["from_cache"] = False  # type: ignore[index]
         return response
 
     async def aclose(self) -> None:

diff --git a/hishel/_async/_transports.py b/hishel/_async/_transports.py
@@ -60,6 +60,7 @@ async def handle_async_request(self, request: Request) -> Response:
             if isinstance(res, httpcore.Response):
                 # Simply use the response if the controller determines it is ready for use.
                 assert isinstance(res.stream, tp.AsyncIterable)
+                res.extensions["from_cache"] = True  # type: ignore[index]
                 return Response(
                     status_code=res.status,
                     headers=res.headers,
@@ -96,6 +97,9 @@ async def handle_async_request(self, request: Request) -> Response:
                 await self._storage.store(key, httpcore_response)
 
                 assert isinstance(httpcore_response.stream, tp.AsyncIterable)
+                httpcore_response.extensions["from_cache"] = (  # type: ignore[index]
+                    httpcore_response.status == 304
+                )
                 return Response(
                     status_code=httpcore_response.status,
                     headers=httpcore_response.headers,
@@ -118,6 +122,7 @@ async def handle_async_request(self, request: Request) -> Response:
             await httpcore_response.aread()
             await self._storage.store(key, httpcore_response)
 
+        response.extensions["from_cache"] = False  # type: ignore[index]
         return response
 
     async def aclose(self) -> None: