Turn MapAsyncIterable into an AsyncGenerator (#199)

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.ExperimentalIncrementalExecutionResults

docs/modules/execution.rst

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

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "graphql-core"
 version = "3.3.0a2"
-description = """
+description = """\
 GraphQL-core is a Python port of GraphQL.js,\
  the JavaScript reference implementation for GraphQL."""
 license = "MIT"

src/graphql/__init__.py

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

src/graphql/execution/__init__.py

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

src/graphql/execution/execute.py

@@ -71,8 +71,7 @@
     is_object_type,
 )
 from .collect_fields import FieldsAndPatches, collect_fields, collect_subfields
-from .flatten_async_iterable import flatten_async_iterable
-from .map_async_iterable import MapAsyncIterable
+from .iterators import flatten_async_iterable, map_async_iterable
 from .middleware import MiddlewareManager
 from .values import get_argument_values, get_directive_values, get_variable_values
 
@@ -1650,7 +1649,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,
@@ -2350,18 +2349,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/flatten_async_iterable.py renamed to src/graphql/execution/iterators.py

@@ -1,4 +1,14 @@
-from typing import AsyncGenerator, AsyncIterable, TypeVar, Union
+from __future__ import annotations  # Python < 3.10
+
+from typing import (
+    Any,
+    AsyncGenerator,
+    AsyncIterable,
+    Awaitable,
+    Callable,
+    TypeVar,
+    Union,
+)
 
 
 try:
@@ -15,10 +25,11 @@ async def aclosing(thing):
 
 
 T = TypeVar("T")
+V = TypeVar("V")
 
 AsyncIterableOrGenerator = Union[AsyncGenerator[T, None], AsyncIterable[T]]
 
-__all__ = ["flatten_async_iterable"]
+__all__ = ["flatten_async_iterable", "map_async_iterable"]
 
 
 async def flatten_async_iterable(
@@ -34,3 +45,23 @@ async def flatten_async_iterable(
             async with aclosing(sub_iterator) as items:  # type: ignore
                 async for item in items:
                     yield item
+
+
+async def map_async_iterable(
+    iterable: AsyncIterable[T], callback: Callable[[T], Awaitable[V]]
+) -> AsyncGenerator[V, None]:
+    """Map an AsyncIterable over a callback function.
+
+    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.
+    """
+
+    aiter = iterable.__aiter__()
+    try:
+        async for element in aiter:
+            yield await callback(element)
+    finally:
+        if hasattr(aiter, "aclose"):
+            await aiter.aclose()

src/graphql/execution/map_async_iterable.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)

tests/execution/test_customize.py

@@ -2,7 +2,7 @@
 
 from pytest import mark, raises
 
-from graphql.execution.flatten_async_iterable import flatten_async_iterable
+from graphql.execution.iterators import flatten_async_iterable
 
 
 try:  # pragma: no cover