refactor: handle athrow edge cases at the same site

Author: winstxnhdw

litestar/response/sse.py

@@ -49,7 +49,7 @@ def __init__(
         if retry_duration is not None:
             chunks.append(f"retry: {retry_duration}\r\n".encode())
 
-        super().__init__(iterable=chunks)
+        super().__init__(iterator=chunks)
 
         if not isinstance(content, (Iterator, AsyncIterator, AsyncIteratorWrapper)) and callable(content):
             content = content()  # type: ignore[unreachable]

litestar/response/streaming.py

@@ -50,7 +50,7 @@ async def async_iterator_to_generator(
 class ASGIStreamingResponse(ASGIResponse):
     """A streaming response."""
 
-    __slots__ = ("_original_generator", "disconnect_event", "iterator")
+    __slots__ = ("disconnect_event", "iterator")
 
     _should_set_content_length = False
 
@@ -93,18 +93,7 @@ def __init__(
         )
 
         self.disconnect_event = Event()
-        self.iterator: AsyncGenerator[str | bytes, None]
-
-        if isinstance(iterator, (AsyncIteratorWrapper, AsyncGenerator)):
-            self.iterator = iterator
-        elif isinstance(iterator, (Iterable, Iterator)):
-            self.iterator = AsyncIteratorWrapper(iterator)
-        elif isinstance(iterator, (AsyncIterable, AsyncIterator)):
-            self.iterator = async_iterator_to_generator(iterator)
-
-        self._original_generator = (
-            iterator.content_async_iterator if hasattr(iterator, "content_async_iterator") else self.iterator  # pyright: ignore[reportAttributeAccessIssue]
-        )
+        self.iterator = iterator if isinstance(iterator, AsyncIterable) else AsyncIteratorWrapper(iterator)
 
     async def _listen_for_disconnect(self, receive: Receive) -> None:
         """Listen for a cancellation message, and if received - call cancel on the cancel scope.
@@ -133,8 +122,11 @@ async def _stream(self, send: Send) -> None:
         async for chunk in self.iterator:
             if self.disconnect_event.is_set():
                 try:
-                    await self.iterator.athrow(ClientDisconnectError)
-                except BaseException:  # noqa: BLE001, S110
+                    if hasattr(self.iterator, "content_async_iterator"):
+                        await self.iterator.content_async_iterator.athrow(ClientDisconnectError)  # pyright: ignore[reportAttributeAccessIssue]
+                    elif isinstance(self.iterator, AsyncGenerator):
+                        await self.iterator.athrow(ClientDisconnectError)
+                except (ClientDisconnectError, StopAsyncIteration):
                     pass
                 finally:
                     break  # noqa: B012

litestar/utils/sync.py

@@ -23,16 +23,16 @@
 S = TypeVar("S", default=None)
 
 
-def iterable_to_generator(iterator: Iterable[T]) -> Generator[T, S, None]:
+def iterable_to_generator(iterable: Iterable[T]) -> Generator[T, S, None]:
     """Convert an iterable to a generator.
 
     Args:
-        iterator: An iterable.
+        iterable: An iterable.
 
     Returns:
         A generator.
     """
-    yield from iterator
+    yield from iterable
 
 
 @overload
@@ -71,22 +71,22 @@ class AsyncIteratorWrapper(AsyncGenerator[T, S]):
 
     __slots__ = ("_original_generator", "generator", "iterator")
 
-    def __init__(self, iterable: Iterable[T]) -> None:
-        """Take a sync iterator or iterable and yields values from it asynchronously.
+    def __init__(self, iterator: Iterable[T]) -> None:
+        """Take a sync iterable and yields values from it asynchronously.
 
         Args:
-            iterable: A sync iterable.
+            iterator: A sync iterable.
         """
         self._original_generator: Generator[T, S, None]
 
-        if isinstance(iterable, Generator):
-            self._original_generator = iterable
-        elif isinstance(iterable, AsyncIteratorWrapper):
-            self._original_generator = iterable._original_generator
+        if isinstance(iterator, Generator):
+            self._original_generator = iterator
+        elif isinstance(iterator, AsyncIteratorWrapper):
+            self._original_generator = iterator._original_generator
         else:
-            self._original_generator = iterable_to_generator(iterable)
+            self._original_generator = iterable_to_generator(iterator)
 
-        self.iterator = iter(iterable)
+        self.iterator = iter(iterator)
         self.generator = self._async_generator()
 
     def _call_next(self) -> T:
@@ -109,10 +109,10 @@ async def __anext__(self) -> T:
         return await self.generator.__anext__()
 
     async def aclose(self) -> None:
-        self._original_generator.close()
+        await sync_to_thread(self._original_generator.close)
 
     async def asend(self, value: S) -> T:
-        return self._original_generator.send(value)
+        return await sync_to_thread(self._original_generator.send, value)
 
     async def athrow(
         self,
@@ -122,9 +122,9 @@ async def athrow(
     ) -> T:
         try:
             return (
-                self._original_generator.throw(typ)
+                await sync_to_thread(self._original_generator.throw, typ)
                 if isinstance(typ, BaseException)
-                else self._original_generator.throw(typ, val, tb)
+                else await sync_to_thread(self._original_generator.throw, typ, val, tb)
             )
         except StopIteration as e:
             raise StopAsyncIteration from e

tests/unit/test_response/test_response_to_asgi_response.py

@@ -63,7 +63,7 @@ def __next__(self) -> str:
 
 class MyAsyncIterator(AsyncIteratorWrapper[str]):
     def __init__(self) -> None:
-        super().__init__(iterable=MySyncIterator())
+        super().__init__(iterator=MySyncIterator())
 
 
 async def test_to_response_returning_litestar_response() -> None: