diff --git a/src/graphql/execution/async_iterables.py b/src/graphql/execution/async_iterables.py index 7646a7da..588f5234 100644 --- a/src/graphql/execution/async_iterables.py +++ b/src/graphql/execution/async_iterables.py @@ -65,6 +65,6 @@ async def map_async_iterable( >>> asyncio.run(doubled_numbers()) [0, 2, 4] """ - async with aclosing(iterable) as items: + async with aclosing(aiter(iterable)) as items: async for item in items: yield await callback(item) diff --git a/tests/execution/test_map_async_iterable.py b/tests/execution/test_map_async_iterable.py index 4133bc6b..9c0a5b2e 100644 --- a/tests/execution/test_map_async_iterable.py +++ b/tests/execution/test_map_async_iterable.py @@ -278,6 +278,44 @@ async def aclose(self): with pytest.raises(StopAsyncIteration): await anext(doubles) + @pytest.mark.parametrize("termination", ["close", "throw", "callback_error"]) + async def closes_iterator_returned_by_async_iterable(termination): + closed = False + + async def source(): + nonlocal closed + try: + yield 1 + finally: + closed = True + + iterator = source() + + class Iterable: + def __aiter__(self): + return iterator + + doubles = map_async_iterable( + Iterable(), throw if termination == "callback_error" else double + ) + try: + if termination == "callback_error": + with pytest.raises(RuntimeError, match="Ouch"): + await anext(doubles) + else: + assert await anext(doubles) == 2 + assert not closed + if termination == "close": + await doubles.aclose() + else: + with pytest.raises(RuntimeError, match="Ouch"): + await doubles.athrow(RuntimeError("Ouch")) + assert closed + with pytest.raises(StopAsyncIteration): + await anext(doubles) + finally: + await iterator.aclose() + async def iterable_exits_on_callback_error(): exited = False diff --git a/tests/execution/test_subscribe.py b/tests/execution/test_subscribe.py index 4e932678..843d787a 100644 --- a/tests/execution/test_subscribe.py +++ b/tests/execution/test_subscribe.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import AsyncIterable, AsyncIterator, Callable +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Callable from typing import ( Any, TypedDict, @@ -1224,6 +1224,45 @@ async def resolve(cls, message, _info) -> str: await subscription.aclose() # type: ignore + async def should_close_iterator_returned_by_async_iterable(): + closed = False + + async def generate_messages(): + nonlocal closed + try: + yield "Hello" + finally: + closed = True + + messages = generate_messages() + + class MessageIterable: + def __aiter__(self): + return messages + + schema = GraphQLSchema( + query=QueryType, + subscription=GraphQLObjectType( + "Subscription", + { + "newMessage": GraphQLField( + GraphQLString, + resolve=lambda message, _info: message, + subscribe=lambda _obj, _info: MessageIterable(), + ) + }, + ), + ) + subscription = subscribe(schema, parse("subscription { newMessage }")) + assert isinstance(subscription, AsyncGenerator) + try: + assert await anext(subscription) == ({"newMessage": "Hello"}, None) + assert not closed + await subscription.aclose() + assert closed + finally: + await messages.aclose() + async def should_close_custom_async_iterator(): class MessageGenerator: closed: bool = False