Use a more tolerant aclosing() context manager

Author: Cito

src/graphql/execution/__init__.py

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

src/graphql/execution/iterators.py renamed to src/graphql/execution/async_iterables.py

@@ -1,5 +1,6 @@
 from __future__ import annotations  # Python < 3.10
 
+from contextlib import AbstractAsyncContextManager
 from typing import (
     Any,
     AsyncGenerator,
@@ -11,25 +12,34 @@
 )
 
 
-try:
-    from contextlib import aclosing
-except ImportError:  # python < 3.10
-    from contextlib import asynccontextmanager
-
-    @asynccontextmanager  # type: ignore
-    async def aclosing(thing):
-        try:
-            yield thing
-        finally:
-            await thing.aclose()
-
+__all__ = ["aclosing", "flatten_async_iterable", "map_async_iterable"]
 
 T = TypeVar("T")
 V = TypeVar("V")
 
 AsyncIterableOrGenerator = Union[AsyncGenerator[T, None], AsyncIterable[T]]
 
-__all__ = ["flatten_async_iterable", "map_async_iterable"]
+
+class aclosing(AbstractAsyncContextManager):
+    """Async context manager for safely finalizing an async iterator or generator.
+
+    Contrary to the function available via the standard library, this one silently
+    ignores the case that custom iterators have no aclose() method.
+    """
+
+    def __init__(self, iterable: AsyncIterableOrGenerator[T]) -> None:
+        self.iterable = iterable
+
+    async def __aenter__(self) -> AsyncIterableOrGenerator[T]:
+        return self.iterable
+
+    async def __aexit__(self, *_exc_info: Any) -> None:
+        try:
+            aclose = self.iterable.aclose  # type: ignore
+        except AttributeError:
+            pass  # do not complain if the iterator has no aclose() method
+        else:
+            await aclose()
 
 
 async def flatten_async_iterable(
@@ -48,7 +58,7 @@ async def flatten_async_iterable(
 
 
 async def map_async_iterable(
-    iterable: AsyncIterable[T], callback: Callable[[T], Awaitable[V]]
+    iterable: AsyncIterableOrGenerator[T], callback: Callable[[T], Awaitable[V]]
 ) -> AsyncGenerator[V, None]:
     """Map an AsyncIterable over a callback function.
 
@@ -58,10 +68,6 @@ async def map_async_iterable(
     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()
+    async with aclosing(iterable) as items:  # type: ignore
+        async for item in items:
+            yield await callback(item)

src/graphql/execution/execute.py

@@ -70,8 +70,8 @@
     is_non_null_type,
     is_object_type,
 )
+from .async_iterables import flatten_async_iterable, map_async_iterable
 from .collect_fields import FieldsAndPatches, collect_fields, collect_subfields
-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
 

tests/execution/test_flatten_async_iterable.py

@@ -2,7 +2,7 @@
 
 from pytest import mark, raises
 
-from graphql.execution.iterators import flatten_async_iterable
+from graphql.execution import flatten_async_iterable
 
 
 try:  # pragma: no cover
@@ -129,3 +129,85 @@ async def nested2() -> AsyncGenerator[float, None]:
         assert await anext(doubles) == 2.2
         with raises(StopAsyncIteration):
             assert await anext(doubles)
+
+    @mark.asyncio
+    async def closes_nested_async_iterators():
+        closed = []
+
+        class Source:
+            def __init__(self):
+                self.counter = 0
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                if self.counter == 2:
+                    raise StopAsyncIteration
+                self.counter += 1
+                return Nested(self.counter)
+
+            async def aclose(self):
+                nonlocal closed
+                closed.append(self.counter)
+
+        class Nested:
+            def __init__(self, value):
+                self.value = value
+                self.counter = 0
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                if self.counter == 2:
+                    raise StopAsyncIteration
+                self.counter += 1
+                return self.value + self.counter / 10
+
+            async def aclose(self):
+                nonlocal closed
+                closed.append(self.value + self.counter / 10)
+
+        doubles = flatten_async_iterable(Source())
+
+        result = [x async for x in doubles]
+
+        assert result == [1.1, 1.2, 2.1, 2.2]
+
+        assert closed == [1.2, 2.2, 2]
+
+    @mark.asyncio
+    async def works_with_nested_async_iterators_that_have_no_close_method():
+        class Source:
+            def __init__(self):
+                self.counter = 0
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                if self.counter == 2:
+                    raise StopAsyncIteration
+                self.counter += 1
+                return Nested(self.counter)
+
+        class Nested:
+            def __init__(self, value):
+                self.value = value
+                self.counter = 0
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                if self.counter == 2:
+                    raise StopAsyncIteration
+                self.counter += 1
+                return self.value + self.counter / 10
+
+        doubles = flatten_async_iterable(Source())
+
+        result = [x async for x in doubles]
+
+        assert result == [1.1, 1.2, 2.1, 2.2]

tests/execution/test_map_async_iterable.py

@@ -3,18 +3,22 @@
 from graphql.execution import map_async_iterable
 
 
-async def map_doubles(x):
+try:  # pragma: no cover
+    anext
+except NameError:  # pragma: no cover (Python < 3.10)
+    # noinspection PyShadowingBuiltins
+    async def anext(iterator):
+        """Return the next item from an async iterator."""
+        return await iterator.__anext__()
+
+
+async def map_doubles(x: int) -> int:
     return x + x
 
 
 def describe_map_async_iterable():
     @mark.asyncio
-    async def test_inner_close_called():
-        """
-        Test that a custom iterator with aclose() gets an aclose() call
-        when outer is closed
-        """
-
+    async def inner_is_closed_when_outer_is_closed():
         class Inner:
             def __init__(self):
                 self.closed = False
@@ -30,19 +34,14 @@ async def __anext__(self):
 
         inner = Inner()
         outer = map_async_iterable(inner, map_doubles)
-        it = outer.__aiter__()
-        assert await it.__anext__() == 2
+        iterator = outer.__aiter__()
+        assert await anext(iterator) == 2
         assert not inner.closed
         await outer.aclose()
         assert inner.closed
 
     @mark.asyncio
-    async def test_inner_close_called_on_callback_err():
-        """
-        Test that a custom iterator with aclose() gets an aclose() call
-        when the callback errors and the outer iterator aborts.
-        """
-
+    async def inner_is_closed_on_callback_error():
         class Inner:
             def __init__(self):
                 self.closed = False
@@ -62,17 +61,11 @@ async def callback(v):
         inner = Inner()
         outer = map_async_iterable(inner, callback)
         with raises(RuntimeError):
-            async for _ in outer:
-                pass
+            await anext(outer)
         assert inner.closed
 
     @mark.asyncio
-    async def test_inner_exit_on_callback_err():
-        """
-        Test that a custom iterator with aclose() gets an aclose() call
-        when the callback errors and the outer iterator aborts.
-        """
-
+    async def test_inner_exits_on_callback_error():
         inner_exit = False
 
         async def inner():
@@ -88,6 +81,35 @@ async def callback(v):
 
         outer = map_async_iterable(inner(), callback)
         with raises(RuntimeError):
-            async for _ in outer:
-                pass
+            await anext(outer)
         assert inner_exit
+
+    @mark.asyncio
+    async def inner_has_no_close_method_when_outer_is_closed():
+        class Inner:
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                return 1
+
+        outer = map_async_iterable(Inner(), map_doubles)
+        iterator = outer.__aiter__()
+        assert await anext(iterator) == 2
+        await outer.aclose()
+
+    @mark.asyncio
+    async def inner_has_no_close_method_on_callback_error():
+        class Inner:
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                return 1
+
+        async def callback(v):
+            raise RuntimeError()
+
+        outer = map_async_iterable(Inner(), callback)
+        with raises(RuntimeError):
+            await anext(outer)