Rewrite MapAsyncIterable using async generator semantics

iterate over and aclose() the iterator

Mapping method must be async

turn MapAsyncIterable into an AsyncGenerator

Author: kristjanvalur

docs/conf.py

@@ -151,7 +151,7 @@
 Middleware
 asyncio.events.AbstractEventLoop
 graphql.execution.collect_fields.FieldsAndPatches
-graphql.execution.map_async_iterable.MapAsyncIterable
+graphql.execution.map_async_iterable.map_async_iterable
 graphql.execution.Middleware
 graphql.execution.execute.DeferredFragmentRecord
 graphql.execution.execute.ExperimentalExecuteMultipleResults

docs/modules/execution.rst

@@ -57,8 +57,6 @@ Execution
 
 .. autofunction:: create_source_event_stream
 
-.. autoclass:: MapAsyncIterable
-
 .. autoclass:: Middleware
 
 .. autoclass:: MiddlewareManager

src/graphql/__init__.py

@@ -439,7 +439,7 @@
     # Subscription
     subscribe,
     create_source_event_stream,
-    MapAsyncIterable,
+    map_async_iterable,
     # Middleware
     Middleware,
     MiddlewareManager,
@@ -707,7 +707,7 @@
     "MiddlewareManager",
     "subscribe",
     "create_source_event_stream",
-    "MapAsyncIterable",
+    "map_async_iterable",
     "validate",
     "ValidationContext",
     "ValidationRule",

src/graphql/execution/__init__.py

@@ -32,7 +32,7 @@
     FormattedIncrementalResult,
     Middleware,
 )
-from .map_async_iterable import MapAsyncIterable
+from .map_async_iterable import map_async_iterable
 from .middleware import MiddlewareManager
 from .values import get_argument_values, get_directive_values, get_variable_values
 
@@ -62,7 +62,7 @@
     "FormattedIncrementalDeferResult",
     "FormattedIncrementalStreamResult",
     "FormattedIncrementalResult",
-    "MapAsyncIterable",
+    "map_async_iterable",
     "Middleware",
     "MiddlewareManager",
     "get_argument_values",

src/graphql/execution/execute.py

@@ -70,7 +70,7 @@
 )
 from .collect_fields import FieldsAndPatches, collect_fields, collect_subfields
 from .flatten_async_iterable import flatten_async_iterable
-from .map_async_iterable import MapAsyncIterable
+from .map_async_iterable import map_async_iterable
 from .middleware import MiddlewareManager
 from .values import get_argument_values, get_directive_values, get_variable_values
 
@@ -1654,7 +1654,7 @@ async def callback(payload: Any) -> AsyncGenerator:
                 await result if isawaitable(result) else result  # type: ignore
             )
 
-        return flatten_async_iterable(MapAsyncIterable(result_or_stream, callback))
+        return flatten_async_iterable(map_async_iterable(result_or_stream, callback))
 
     def execute_deferred_fragment(
         self,
@@ -2319,18 +2319,20 @@ def subscribe(
     if isinstance(result, ExecutionResult):
         return result
     if isinstance(result, AsyncIterable):
-        return MapAsyncIterable(result, ensure_single_execution_result)
+        return map_async_iterable(result, ensure_single_execution_result)
 
     async def await_result() -> Union[AsyncIterator[ExecutionResult], ExecutionResult]:
         result_or_iterable = await result  # type: ignore
         if isinstance(result_or_iterable, AsyncIterable):
-            return MapAsyncIterable(result_or_iterable, ensure_single_execution_result)
+            return map_async_iterable(
+                result_or_iterable, ensure_single_execution_result
+            )
         return result_or_iterable
 
     return await_result()
 
 
-def ensure_single_execution_result(
+async def ensure_single_execution_result(
     result: Union[
         ExecutionResult,
         InitialIncrementalExecutionResult,

src/graphql/execution/map_async_iterable.py

@@ -1,118 +1,26 @@
 from __future__ import annotations  # Python < 3.10
 
-from asyncio import CancelledError, Event, Task, ensure_future, wait
-from concurrent.futures import FIRST_COMPLETED
-from inspect import isasyncgen, isawaitable
-from types import TracebackType
-from typing import Any, AsyncIterable, Callable, Optional, Set, Type, Union, cast
+from typing import Any, AsyncIterable, Awaitable, Callable
 
 
-__all__ = ["MapAsyncIterable"]
+__all__ = ["map_async_iterable"]
 
 
-# noinspection PyAttributeOutsideInit
-class MapAsyncIterable:
+async def map_async_iterable(
+    iterable: AsyncIterable[Any], callback: Callable[[Any], Awaitable[Any]]
+) -> None:
     """Map an AsyncIterable over a callback function.
 
-    Given an AsyncIterable and a callback function, return an AsyncIterator which
-    produces values mapped via calling the callback function.
-
-    When the resulting AsyncIterator is closed, the underlying AsyncIterable will also
-    be closed.
+    Given an AsyncIterable and an async callback callable, return an AsyncGenerator
+    which produces values mapped via calling the callback.
+    If the inner iterator supports an `aclose()` method, it will be called when
+    the generator finishes or closes.
     """
 
-    def __init__(self, iterable: AsyncIterable, callback: Callable) -> None:
-        self.iterator = iterable.__aiter__()
-        self.callback = callback
-        self._close_event = Event()
-
-    def __aiter__(self) -> MapAsyncIterable:
-        """Get the iterator object."""
-        return self
-
-    async def __anext__(self) -> Any:
-        """Get the next value of the iterator."""
-        if self.is_closed:
-            if not isasyncgen(self.iterator):
-                raise StopAsyncIteration
-            value = await self.iterator.__anext__()
-        else:
-            aclose = ensure_future(self._close_event.wait())
-            anext = ensure_future(self.iterator.__anext__())
-
-            try:
-                pending: Set[Task] = (
-                    await wait([aclose, anext], return_when=FIRST_COMPLETED)
-                )[1]
-            except CancelledError:
-                # cancel underlying tasks and close
-                aclose.cancel()
-                anext.cancel()
-                await self.aclose()
-                raise  # re-raise the cancellation
-
-            for task in pending:
-                task.cancel()
-
-            if aclose.done():
-                raise StopAsyncIteration
-
-            error = anext.exception()
-            if error:
-                raise error
-
-            value = anext.result()
-
-        result = self.callback(value)
-
-        return await result if isawaitable(result) else result
-
-    async def athrow(
-        self,
-        type_: Union[BaseException, Type[BaseException]],
-        value: Optional[BaseException] = None,
-        traceback: Optional[TracebackType] = None,
-    ) -> None:
-        """Throw an exception into the asynchronous iterator."""
-        if self.is_closed:
-            return
-        athrow = getattr(self.iterator, "athrow", None)
-        if athrow:
-            await athrow(type_, value, traceback)
-        else:
-            await self.aclose()
-            if value is None:
-                if traceback is None:
-                    raise type_
-                value = (
-                    type_
-                    if isinstance(value, BaseException)
-                    else cast(Type[BaseException], type_)()
-                )
-            if traceback is not None:
-                value = value.with_traceback(traceback)
-            raise value
-
-    async def aclose(self) -> None:
-        """Close the iterator."""
-        if not self.is_closed:
-            aclose = getattr(self.iterator, "aclose", None)
-            if aclose:
-                try:
-                    await aclose()
-                except RuntimeError:
-                    pass
-            self.is_closed = True
-
-    @property
-    def is_closed(self) -> bool:
-        """Check whether the iterator is closed."""
-        return self._close_event.is_set()
-
-    @is_closed.setter
-    def is_closed(self, value: bool) -> None:
-        """Mark the iterator as closed."""
-        if value:
-            self._close_event.set()
-        else:
-            self._close_event.clear()
+    aiter = iterable.__aiter__()
+    try:
+        async for element in aiter:
+            yield await callback(element)
+    finally:
+        if hasattr(aiter, "aclose"):
+            await aiter.aclose()

tests/execution/test_customize.py

@@ -1,6 +1,8 @@
+from inspect import isasyncgen
+
 from pytest import mark
 
-from graphql.execution import ExecutionContext, MapAsyncIterable, execute, subscribe
+from graphql.execution import ExecutionContext, execute, subscribe
 from graphql.language import parse
 from graphql.type import GraphQLField, GraphQLObjectType, GraphQLSchema, GraphQLString
 
@@ -77,7 +79,7 @@ async def custom_foo():
             root_value=Root(),
             subscribe_field_resolver=lambda root, _info: root.custom_foo(),
         )
-        assert isinstance(subscription, MapAsyncIterable)
+        assert isasyncgen(subscription)
 
         assert await anext(subscription) == (
             {"foo": "FooValue"},
@@ -121,6 +123,6 @@ def resolve_foo(message, _info):
             context_value={},
             execution_context_class=TestExecutionContext,
         )
-        assert isinstance(subscription, MapAsyncIterable)
+        assert isasyncgen(subscription)
 
         assert await anext(subscription) == ({"foo": "bar"}, None)