Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use Iterable instead Iterator on iterate_in_threadpool #2362

Merged
merged 7 commits into from
Dec 20, 2023
Merged
5 changes: 3 additions & 2 deletions starlette/concurrency.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,10 +55,11 @@ def _next(iterator: typing.Iterator[T]) -> T:


async def iterate_in_threadpool(
iterator: typing.Iterator[T],
iterator: typing.Iterable[T],
) -> typing.AsyncIterator[T]:
as_iterator = iter(iterator)
while True:
try:
yield await anyio.to_thread.run_sync(_next, iterator)
yield await anyio.to_thread.run_sync(_next, as_iterator)
except _StopIteration:
break
2 changes: 1 addition & 1 deletion starlette/responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ def __init__(


Content = typing.Union[str, bytes]
SyncContentStream = typing.Iterator[Content]
SyncContentStream = typing.Iterable[Content]
AsyncContentStream = typing.AsyncIterable[Content]
ContentStream = typing.Union[AsyncContentStream, SyncContentStream]

Expand Down
11 changes: 10 additions & 1 deletion tests/test_concurrency.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import pytest

from starlette.applications import Starlette
from starlette.concurrency import run_until_first_complete
from starlette.concurrency import iterate_in_threadpool, run_until_first_complete
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Route
Expand Down Expand Up @@ -40,3 +40,12 @@ def endpoint(request: Request) -> Response:

resp = client.get("/")
assert resp.content == b"data"


@pytest.mark.anyio
async def test_iterate_in_threadpool() -> None:
class CustomIterable:
def __iter__(self):
yield from range(3)

assert [v async for v in iterate_in_threadpool(CustomIterable())] == [0, 1, 2]