Support custom async iterables

Add is_async_iterable parameter for custom iterables (#258)

Author: Cito

docs/modules/pyutils.rst

@@ -16,6 +16,7 @@ PyUtils
 .. autofunction:: identity_func
 .. autofunction:: inspect
 .. autofunction:: is_awaitable
+.. autofunction:: is_async_iterable
 .. autofunction:: is_collection
 .. autofunction:: is_iterable
 .. autofunction:: natural_comparison_key

src/graphql/execution/execute.py

@@ -47,7 +47,12 @@
     inspect,
     is_iterable,
 )
-from ..pyutils.is_awaitable import is_awaitable as default_is_awaitable
+from ..pyutils.is_awaitable import (
+    is_async_iterable as default_is_async_iterable,
+)
+from ..pyutils.is_awaitable import (
+    is_awaitable as default_is_awaitable,
+)
 from ..type import (
     GraphQLAbstractType,
     GraphQLField,
@@ -201,6 +206,9 @@ class ExecutionContext(IncrementalPublisherContext):
     is_awaitable: Callable[[Any], TypeGuard[Awaitable]] = staticmethod(
         default_is_awaitable  # type: ignore
     )
+    is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] = staticmethod(
+        default_is_async_iterable  # type: ignore
+    )
 
     def __init__(
         self,
@@ -216,6 +224,7 @@ def __init__(
         enable_early_execution: bool = False,
         middleware_manager: MiddlewareManager | None = None,
         is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None,
+        is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None,
     ) -> None:
         self.schema = schema
         self.fragments = fragments
@@ -229,6 +238,7 @@ def __init__(
         self.enable_early_execution = enable_early_execution
         self.middleware_manager = middleware_manager
         self.is_awaitable = is_awaitable or default_is_awaitable
+        self.is_async_iterable = is_async_iterable or default_is_async_iterable
         self.errors = None
         self.cancellable_streams = None
         self._canceled_iterators: set[AsyncIterator] = set()
@@ -251,6 +261,7 @@ def build(
         enable_early_execution: bool = False,
         middleware: Middleware | None = None,
         is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None,
+        is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None,
         **custom_args: Any,
     ) -> list[GraphQLError] | ExecutionContext:
         """Build an execution context
@@ -325,6 +336,7 @@ def build(
             enable_early_execution,
             middleware_manager,
             is_awaitable,
+            is_async_iterable,
             **custom_args,
         )
 
@@ -1041,7 +1053,7 @@ def complete_list_value(
         """
         item_type = return_type.of_type
 
-        if isinstance(result, AsyncIterable):
+        if self.is_async_iterable(result):
             async_iterator = result.__aiter__()
 
             return self.complete_async_iterator_value(
@@ -1582,8 +1594,8 @@ def map_source_to_response(
         as it is nearly identical to the "ExecuteQuery" algorithm,
         for which :func:`~graphql.execution.execute` is also used.
         """
-        if not isinstance(result_or_stream, AsyncIterable):
-            return result_or_stream  # pragma: no cover
+        if not self.is_async_iterable(result_or_stream):
+            return cast("ExecutionResult", result_or_stream)  # pragma: no cover
 
         build_context = self.build_per_event_execution_context
 
@@ -2103,6 +2115,7 @@ def execute(
     middleware: Middleware | None = None,
     execution_context_class: type[ExecutionContext] | None = None,
     is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None,
+    is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None,
     **custom_context_args: Any,
 ) -> AwaitableOrValue[ExecutionResult]:
     """Execute a GraphQL operation.
@@ -2137,6 +2150,7 @@ def execute(
         middleware,
         execution_context_class,
         is_awaitable,
+        is_async_iterable,
         **custom_context_args,
     )
     if isinstance(result, ExecutionResult):
@@ -2167,6 +2181,7 @@ def experimental_execute_incrementally(
     middleware: Middleware | None = None,
     execution_context_class: type[ExecutionContext] | None = None,
     is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None,
+    is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None,
     **custom_context_args: Any,
 ) -> AwaitableOrValue[ExecutionResult | ExperimentalIncrementalExecutionResults]:
     """Execute GraphQL operation incrementally (internal implementation).
@@ -2197,6 +2212,7 @@ def experimental_execute_incrementally(
         enable_early_execution,
         middleware,
         is_awaitable,
+        is_async_iterable,
         **custom_context_args,
     )
 
@@ -2692,7 +2708,7 @@ def assert_event_stream(result: Any) -> AsyncIterable:
         raise result
 
     # Assert field returned an event stream, otherwise yield an error.
-    if not isinstance(result, AsyncIterable):
+    if not default_is_async_iterable(result):
         msg = (
             "Subscription field must return AsyncIterable."
             f" Received: {inspect(result)}."

src/graphql/graphql.py

@@ -17,7 +17,7 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import Awaitable, Callable
+    from collections.abc import AsyncIterable, Awaitable, Callable
     from typing import TypeGuard
 
     from .pyutils import AwaitableOrValue
@@ -37,6 +37,7 @@ async def graphql(
     middleware: Middleware | None = None,
     execution_context_class: type[ExecutionContext] | None = None,
     is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None,
+    is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None,
 ) -> ExecutionResult:
     """Execute a GraphQL operation asynchronously.
 
@@ -84,6 +85,8 @@ async def graphql(
       The execution context class to use to build the context
     :arg is_awaitable:
       The predicate to be used for checking whether values are awaitable
+    :arg is_async_iterable:
+      The predicate to be used for checking whether values are async iterables
     """
     # Always return asynchronously for a consistent API.
     result = graphql_impl(
@@ -98,6 +101,7 @@ async def graphql(
         middleware,
         execution_context_class,
         is_awaitable,
+        is_async_iterable,
     )
 
     if default_is_awaitable(result):
@@ -111,6 +115,11 @@ def assume_not_awaitable(_value: Any) -> TypeGuard[Awaitable]:
     return False
 
 
+def assume_not_async_iterable(_value: Any) -> TypeGuard[AsyncIterable]:
+    """Replacement for is_async_iterable if everything is assumed to be synchronous."""
+    return False
+
+
 def graphql_sync(
     schema: GraphQLSchema,
     source: str | Source,
@@ -138,6 +147,7 @@ def graphql_sync(
         if callable(check_sync)
         else (None if check_sync else assume_not_awaitable)
     )
+    is_async_iterable = assume_not_async_iterable if not check_sync else None
     result = graphql_impl(
         schema,
         source,
@@ -150,6 +160,7 @@ def graphql_sync(
         middleware,
         execution_context_class,
         is_awaitable,
+        is_async_iterable,
     )
 
     # Assert that the execution was synchronous.
@@ -173,6 +184,7 @@ def graphql_impl(
     middleware: Middleware | None,
     execution_context_class: type[ExecutionContext] | None,
     is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None,
+    is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None,
 ) -> AwaitableOrValue[ExecutionResult]:
     """Execute a query, return asynchronously only if necessary."""
     # Validate Schema
@@ -206,4 +218,5 @@ def graphql_impl(
         middleware,
         execution_context_class,
         is_awaitable,
+        is_async_iterable,
     )

src/graphql/pyutils/__init__.py

@@ -24,7 +24,7 @@
 from .group_by import group_by
 from .identity_func import identity_func
 from .inspect import inspect
-from .is_awaitable import is_awaitable
+from .is_awaitable import is_awaitable, is_async_iterable
 from .is_iterable import is_collection, is_iterable
 from .natural_compare import natural_comparison_key
 from .awaitable_or_value import AwaitableOrValue
@@ -59,6 +59,7 @@
     "group_by",
     "identity_func",
     "inspect",
+    "is_async_iterable",
     "is_awaitable",
     "is_collection",
     "is_description",

src/graphql/pyutils/is_awaitable.py

@@ -7,9 +7,9 @@
 from typing import TYPE_CHECKING, Any, TypeGuard
 
 if TYPE_CHECKING:
-    from collections.abc import Awaitable
+    from collections.abc import AsyncIterable, Awaitable
 
-__all__ = ["is_awaitable"]
+__all__ = ["is_async_iterable", "is_awaitable"]
 
 CO_ITERABLE_COROUTINE = inspect.CO_ITERABLE_COROUTINE
 
@@ -31,3 +31,12 @@ def is_awaitable(value: Any) -> TypeGuard[Awaitable]:
         # check for other awaitables (e.g. futures)
         or hasattr(value, "__await__")
     )
+
+
+def is_async_iterable(value: Any) -> TypeGuard[AsyncIterable]:
+    """Return True if object is an asynchronous iterable.
+
+    Instead of testing whether the object is an instance of abc.AsyncIterable, we
+    check the existence of an `__aiter__` attribute. This is much faster.
+    """
+    return hasattr(value, "__aiter__")

tests/execution/test_lists.py

@@ -1,5 +1,5 @@
-from collections.abc import AsyncGenerator
-from typing import Any
+from collections.abc import AsyncGenerator, AsyncIterable, Callable
+from typing import Any, TypeGuard
 
 import pytest
 
@@ -204,6 +204,62 @@ async def __anext__(self):
             None,
         )
 
+    @pytest.mark.filterwarnings("ignore:.* was never awaited:RuntimeWarning")
+    async def can_customize_detection_of_async_iterables():
+        class CustomIterable:
+            """An object that is both an iterable and an async iterable."""
+
+            stop = False
+
+            def __iter__(self):
+                return self
+
+            def __next__(self):
+                if self.stop:
+                    raise StopIteration
+                self.stop = True
+                return "hello"
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                if self.stop:
+                    raise StopAsyncIteration
+                self.stop = True
+                return "world"
+
+        assert await _complete(CustomIterable()) == (
+            {"listField": ["world"]},
+            None,
+        )
+
+        async def _complete_custom(
+            is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None,
+        ):
+            return execute(
+                build_schema("type Query { listField: [String] }"),
+                parse("{ listField }"),
+                Data(CustomIterable()),
+                is_async_iterable=is_async_iterable,
+            )
+
+        def use_async(result: Any) -> TypeGuard[AsyncIterable]:
+            return isinstance(result, AsyncIterable)
+
+        result = await _complete_custom(use_async)
+
+        assert is_awaitable(result)
+        assert await result == ({"listField": ["world"]}, None)
+
+        def use_sync(result: Any) -> TypeGuard[AsyncIterable]:
+            return isinstance(result, AsyncIterable) and not hasattr(result, "__iter__")
+
+        result = await _complete_custom(use_sync)
+
+        assert not is_awaitable(result)
+        assert result == ({"listField": ["hello"]}, None)
+
     async def handles_an_async_generator_that_throws():
         async def list_field():
             yield "two"