Move both async generatrs into common iterators module

Author: kristjanvalur

src/graphql/execution/__init__.py

@@ -32,7 +32,7 @@
     FormattedIncrementalResult,
     Middleware,
 )
-from .map_async_iterable import map_async_iterable
+from .iterators import map_async_iterable
 from .middleware import MiddlewareManager
 from .values import get_argument_values, get_directive_values, get_variable_values
 

src/graphql/execution/execute.py

@@ -69,8 +69,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 map_async_iterable
+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
 

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:
@@ -18,7 +28,7 @@ async def aclosing(thing):
 
 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 +44,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[Any]]
+) -> 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

@@ -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