fix(streaming): open streams with client.stream so httpx2 type-checks

Author: tharropoulos

src/typesense/async_/api_call.py

@@ -33,6 +33,7 @@
 
 import asyncio
 import sys
+from contextlib import AsyncExitStack
 from types import MappingProxyType, TracebackType
 
 import httpx
@@ -54,6 +55,7 @@
 from typesense.http_backend import (
     ASYNC_CLIENT_TYPES,
     AsyncClientType,
+    ResponseType,
     backend_errors,
     verify_option,
 )
@@ -648,38 +650,41 @@ async def _open_stream(
         headers = request_kwargs.get("headers", {})
         headers["Accept"] = "text/event-stream"
         timeout = self._client.timeout
-        request = self._client.build_request(
-            method,
-            url,
-            params=typing.cast(
-                typing.Optional[_QueryParams],
-                request_kwargs.get("params"),
-            ),
-            content=request_kwargs.get("content"),
-            headers=headers,
-            timeout=(
-                timeout.connect,
-                self.config.stream_read_timeout_seconds,
-                timeout.write,
-                timeout.pool,
-            ),
+        # Annotated so httpx and httpx2 responses unify as ``ResponseType``.
+        response_context: typing.AsyncContextManager[ResponseType] = (
+            self._client.stream(
+                method,
+                url,
+                params=typing.cast(
+                    typing.Optional[_QueryParams],
+                    request_kwargs.get("params"),
+                ),
+                content=request_kwargs.get("content"),
+                headers=headers,
+                timeout=(
+                    timeout.connect,
+                    self.config.stream_read_timeout_seconds,
+                    timeout.write,
+                    timeout.pool,
+                ),
+            )
         )
 
+        # Owns the concurrency slot and the response until the stream is closed.
+        exit_stack = AsyncExitStack()
         await self._concurrency_limit.acquire()
+        exit_stack.callback(self._concurrency_limit.release)
         try:
-            response = await self._client.send(request, stream=True)
+            response = await exit_stack.enter_async_context(response_context)
             if response.status_code < 200 or response.status_code >= 300:
-                try:
-                    await response.aread()
-                finally:
-                    await response.aclose()
+                await response.aread()
                 self.request_handler.raise_for_status(response)
         except BaseException:
-            self._concurrency_limit.release()
+            await exit_stack.aclose()
             raise
 
         self.node_manager.set_node_health(node, is_healthy=True)
-        return AsyncSearchStream(response, self._concurrency_limit.release)
+        return AsyncSearchStream(response, exit_stack)
 
     def _prepare_request_params(
         self,

src/typesense/async_/stream.py

@@ -14,6 +14,7 @@
 """
 
 import sys
+from contextlib import AsyncExitStack
 from types import TracebackType
 
 from typesense.exceptions import TypesenseClientError
@@ -46,19 +47,19 @@ class AsyncSearchStream(typing.Generic[TFinal]):
     def __init__(
         self,
         response: ResponseType,
-        on_close: typing.Callable[[], None],
+        exit_stack: AsyncExitStack,
     ) -> None:
         """
         Initialize the stream.
 
         Args:
             response (httpx.Response | httpx2.Response): A successful response
                 opened with ``stream=True``.
-            on_close (Callable[[], None]): Called once when the stream is closed,
-                to release the request's concurrency slot.
+            exit_stack (AsyncExitStack): Closes the response and releases the
+                request's concurrency slot when the stream is closed.
         """
         self.response = response
-        self._on_close = on_close
+        self._exit_stack = exit_stack
         self._closed = False
         self._final: typing.Optional[TFinal] = None
         self._decoder = SSEDecoder()
@@ -112,14 +113,11 @@ async def aclose(self) -> None:
         await self._close_response()
 
     async def _close_response(self) -> None:
-        """Close the response once, then run ``on_close``."""
+        """Close the response and release its concurrency slot, once."""
         if self._closed:
             return
         self._closed = True
-        try:
-            await self.response.aclose()
-        finally:
-            self._on_close()
+        await self._exit_stack.aclose()
 
     async def _iter_chunks(self) -> typing.AsyncGenerator[MessageChunk, None]:
         """Yield the answer pieces and keep the final search response."""

src/typesense/sync/api_call.py

@@ -33,6 +33,7 @@
 
 import time
 import sys
+from contextlib import ExitStack
 from types import MappingProxyType, TracebackType
 
 import httpx
@@ -54,6 +55,7 @@
 from typesense.http_backend import (
     CLIENT_TYPES,
     SyncClientType,
+    ResponseType,
     backend_errors,
     verify_option,
 )
@@ -648,38 +650,41 @@ def _open_stream(
         headers = request_kwargs.get("headers", {})
         headers["Accept"] = "text/event-stream"
         timeout = self._client.timeout
-        request = self._client.build_request(
-            method,
-            url,
-            params=typing.cast(
-                typing.Optional[_QueryParams],
-                request_kwargs.get("params"),
-            ),
-            content=request_kwargs.get("content"),
-            headers=headers,
-            timeout=(
-                timeout.connect,
-                self.config.stream_read_timeout_seconds,
-                timeout.write,
-                timeout.pool,
-            ),
+        # Annotated so httpx and httpx2 responses unify as ``ResponseType``.
+        response_context: typing.ContextManager[ResponseType] = (
+            self._client.stream(
+                method,
+                url,
+                params=typing.cast(
+                    typing.Optional[_QueryParams],
+                    request_kwargs.get("params"),
+                ),
+                content=request_kwargs.get("content"),
+                headers=headers,
+                timeout=(
+                    timeout.connect,
+                    self.config.stream_read_timeout_seconds,
+                    timeout.write,
+                    timeout.pool,
+                ),
+            )
         )
 
+        # Owns the concurrency slot and the response until the stream is closed.
+        exit_stack = ExitStack()
         self._concurrency_limit.acquire()
+        exit_stack.callback(self._concurrency_limit.release)
         try:
-            response = self._client.send(request, stream=True)
+            response = exit_stack.enter_context(response_context)
             if response.status_code < 200 or response.status_code >= 300:
-                try:
-                    response.read()
-                finally:
-                    response.close()
+                response.read()
                 self.request_handler.raise_for_status(response)
         except BaseException:
-            self._concurrency_limit.release()
+            exit_stack.close()
             raise
 
         self.node_manager.set_node_health(node, is_healthy=True)
-        return SearchStream(response, self._concurrency_limit.release)
+        return SearchStream(response, exit_stack)
 
     def _prepare_request_params(
         self,

src/typesense/sync/stream.py

@@ -14,6 +14,7 @@
 """
 
 import sys
+from contextlib import ExitStack
 from types import TracebackType
 
 from typesense.exceptions import TypesenseClientError
@@ -46,19 +47,19 @@ class SearchStream(typing.Generic[TFinal]):
     def __init__(
         self,
         response: ResponseType,
-        on_close: typing.Callable[[], None],
+        exit_stack: ExitStack,
     ) -> None:
         """
         Initialize the stream.
 
         Args:
             response (httpx.Response | httpx2.Response): A successful response
                 opened with ``stream=True``.
-            on_close (Callable[[], None]): Called once when the stream is closed,
-                to release the request's concurrency slot.
+            exit_stack (ExitStack): Closes the response and releases the
+                request's concurrency slot when the stream is closed.
         """
         self.response = response
-        self._on_close = on_close
+        self._exit_stack = exit_stack
         self._closed = False
         self._final: typing.Optional[TFinal] = None
         self._decoder = SSEDecoder()
@@ -112,14 +113,11 @@ def close(self) -> None:
         self._close_response()
 
     def _close_response(self) -> None:
-        """Close the response once, then run ``on_close``."""
+        """Close the response and release its concurrency slot, once."""
         if self._closed:
             return
         self._closed = True
-        try:
-            self.response.close()
-        finally:
-            self._on_close()
+        self._exit_stack.close()
 
     def _iter_chunks(self) -> typing.Generator[MessageChunk, None, None]:
         """Yield the answer pieces and keep the final search response."""

utils/run-unasync.py

@@ -35,6 +35,9 @@ def collect_class_replacements(source_dir: Path) -> dict[str, str]:
     # Defined in the shared ``typesense.sse`` module, outside async_.
     replacements["aiter_events"] = "iter_events"
     replacements["aiter_bytes"] = "iter_bytes"
+    replacements["AsyncExitStack"] = "ExitStack"
+    replacements["AsyncContextManager"] = "ContextManager"
+    replacements["enter_async_context"] = "enter_context"
     # ``AsyncGenerator`` takes two type arguments, but ``Generator`` needs three
     # before Python 3.13.
     replacements["Generator[MessageChunk, None]"] = "Generator[MessageChunk, None, None]"