Get rid of ext close callback, and other tweaks

johtso · Jun 23, 2021 · 3c41933 · 3c41933
1 parent 9c4659e
commit 3c41933
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 20 deletions.
diff --git a/httpx_caching/_async/_transport.py b/httpx_caching/_async/_transport.py
@@ -120,9 +120,8 @@ async def _io_close_response_stream(
     def wrap_response_stream(
         self, key: str, response: Response, vary_header_values: dict
     ) -> Response:
-        wrapped_stream = ByteStreamWrapper(
-            response.stream, response.extensions.get("aclose")
-        )
+        response_stream: AsyncByteStream = response.stream  # type: ignore
+        wrapped_stream = ByteStreamWrapper(response_stream)
         response.stream = wrapped_stream
 
         async def callback(response_body: bytes):

diff --git a/httpx_caching/_models.py b/httpx_caching/_models.py
@@ -1,7 +1,9 @@
 import dataclasses
-from typing import AsyncIterable, Iterable, Union
+from typing import Union
 
-from httpx import Headers
+from httpx import ByteStream, Headers
+
+from ._utils import ByteStreamWrapper
 
 
 @dataclasses.dataclass
@@ -12,7 +14,7 @@ class Response:
 
     status_code: int
     headers: Headers
-    stream: Union[Iterable[bytes], AsyncIterable[bytes]]
+    stream: Union[ByteStream, ByteStreamWrapper]
     extensions: dict = dataclasses.field(default_factory=dict)
 
     @classmethod

diff --git a/httpx_caching/_sync/_transport.py b/httpx_caching/_sync/_transport.py
@@ -118,9 +118,8 @@ def _io_close_response_stream(self, action: protocol.CloseResponseStream) -> Non
     def wrap_response_stream(
         self, key: str, response: Response, vary_header_values: dict
     ) -> Response:
-        wrapped_stream = ByteStreamWrapper(
-            response.stream, response.extensions.get("close")
-        )
+        response_stream: SyncByteStream = response.stream  # type: ignore
+        wrapped_stream = ByteStreamWrapper(response_stream)
         response.stream = wrapped_stream
 
         def callback(response_body: bytes):

diff --git a/httpx_caching/_utils.py b/httpx_caching/_utils.py
@@ -1,11 +1,9 @@
 import threading
 from typing import (
-    AsyncIterable,
     AsyncIterator,
     Awaitable,
     Callable,
     Generator,
-    Iterable,
     Iterator,
     Optional,
     Tuple,
@@ -23,8 +21,7 @@
 class ByteStreamWrapper:
     def __init__(
         self,
-        stream: Union[Iterable[bytes], AsyncIterable[bytes]],
-        stream_close: Optional[Callable],
+        stream: Union[httpx.SyncByteStream, httpx.AsyncByteStream],
         callback: Optional[Callable] = None,
     ) -> None:
         """
@@ -33,7 +30,6 @@ def __init__(
         """
         self.stream = stream
         self.callback = callback or (lambda *args, **kwargs: None)
-        self.stream_close = stream_close
 
         self.buffer = bytearray()
         self.callback_called = False
@@ -43,7 +39,7 @@ def _on_read_finish(self):
             self.callback(bytes(self.buffer))
             self.callback_called = True
 
-    async def a_on_read_finish(self):
+    async def _a_on_read_finish(self):
         if not self.callback_called:
             await self.callback(bytes(self.buffer))
             self.callback_called = True
@@ -58,15 +54,13 @@ async def __aiter__(self) -> AsyncIterator[bytes]:
         async for chunk in self.stream:  # type: ignore
             self.buffer.extend(chunk)
             yield chunk
-        await self.a_on_read_finish()
+        await self._a_on_read_finish()
 
     def close(self) -> None:
-        if self.stream_close:
-            self.stream_close()  # type: ignore
+        self.stream.close()  # type: ignore
 
     async def aclose(self) -> None:
-        if self.stream_close:
-            await self.stream_close()  # type: ignore
+        await self.stream.aclose()  # type: ignore
 
 
 YieldType = TypeVar("YieldType")