polish(incremental): simplify iterator

Replicates graphql/graphql-js@5cd5001

Author: Cito

src/graphql/execution/incremental_graph.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 from asyncio import (
-    CancelledError,
     Future,
     Task,
     ensure_future,
@@ -25,7 +24,7 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Awaitable, Generator, Iterable, Sequence
+    from collections.abc import Awaitable, Generator, Iterable, Sequence
 
     from ..error.graphql_error import GraphQLError
     from .types import (
@@ -48,7 +47,7 @@ class IncrementalGraph:
 
     _root_nodes: dict[SubsequentResultRecord, None]
     _completed_queue: list[IncrementalDataRecordResult]
-    _next_queue: list[Future[Iterable[IncrementalDataRecordResult]]]
+    _next_queue: list[Future[Iterable[IncrementalDataRecordResult] | None]]
 
     _tasks: set[Task[Any]]
 
@@ -87,24 +86,31 @@ def add_completed_reconcilable_deferred_grouped_field_set(
                 incremental_data_records, deferred_records
             )
 
-    async def completed_incremental_data(
+    def current_completed_batch(
         self,
-    ) -> AsyncGenerator[Iterable[IncrementalDataRecordResult], None]:
-        """Asynchronously yield completed incremental data record results."""
+    ) -> Generator[IncrementalDataRecordResult, None, None]:
+        """Yield the current completed batch of incremental data record results."""
+        queue = self._completed_queue
+        while queue:
+            yield queue.pop(0)
+        if not self._root_nodes:
+            self.abort()
+
+    def next_completed_batch(
+        self,
+    ) -> Future[Iterable[IncrementalDataRecordResult] | None]:
+        """Return a future that resolves to the next completed batch."""
         loop = get_running_loop()
-        while True:
-            if self._completed_queue:
-                first_result = self._completed_queue.pop(0)
-                yield self._yield_current_completed_incremental_data(first_result)
-            else:
-                future: Future[Iterable[IncrementalDataRecordResult]] = (
-                    loop.create_future()
-                )
-                self._next_queue.append(future)
-                try:
-                    yield await future
-                except CancelledError:
-                    break  # pragma: no cover
+        future: Future[Iterable[IncrementalDataRecordResult] | None] = (
+            loop.create_future()
+        )
+        self._next_queue.append(future)
+        return future
+
+    def abort(self) -> None:
+        """Abort the incremental graph execution."""
+        for resolve in self._next_queue:
+            resolve.set_result(None)
 
     def has_next(self) -> bool:
         """Check if there are more results to process."""
@@ -332,11 +338,7 @@ def _yield_current_completed_incremental_data(
     ) -> Generator[IncrementalDataRecordResult, None, None]:
         """Yield the current completed incremental data."""
         yield first_result
-        queue = self._completed_queue
-        while queue:
-            yield queue.pop(0)
-        if not self._root_nodes:
-            self.stop_incremental_data()
+        yield from self.current_completed_batch()
 
     def _enqueue(self, completed: IncrementalDataRecordResult) -> None:
         """Enqueue completed incremental data record result."""

src/graphql/execution/incremental_publisher.py

@@ -27,7 +27,7 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Sequence
+    from collections.abc import AsyncGenerator, Iterable, Sequence
 
     from ..error import GraphQLError
     from .types import (
@@ -134,38 +134,33 @@ async def _subscribe(
         incremental_graph = self._incremental_graph
         check_has_next = incremental_graph.has_next
         handle_completed_incremental_data = self._handle_completed_incremental_data
-        completed_incremental_data = incremental_graph.completed_incremental_data()
-        # use the raw iterator rather than 'async for' so as not to end the iterator
-        # when exiting the loop with the next value
-        get_next_results = completed_incremental_data.__aiter__().__anext__
-        is_done = False
-        try:
-            while not is_done:
-                try:
-                    completed_results = await get_next_results()
-                except StopAsyncIteration:  # pragma: no cover
-                    break
 
-                context = SubsequentIncrementalExecutionResultContext([], [], [])
-                for completed_result in completed_results:
-                    await handle_completed_incremental_data(completed_result, context)
+        try:
+            while True:
+                batch: Iterable[IncrementalDataRecordResult] | None = (
+                    incremental_graph.current_completed_batch()
+                )
 
-                if context.incremental or context.completed:  # pragma: no branch
-                    has_next = check_has_next()
+                while batch is not None:
+                    context = SubsequentIncrementalExecutionResultContext([], [], [])
+                    for completed_result in batch:
+                        await handle_completed_incremental_data(
+                            completed_result, context
+                        )
 
-                    if not has_next:
-                        is_done = True
+                    if context.incremental or context.completed:  # pragma: no branch
+                        has_next = check_has_next()
 
-                    subsequent_incremental_execution_result = (
-                        SubsequentIncrementalExecutionResult(
+                        yield SubsequentIncrementalExecutionResult(
                             has_next=has_next,
                             pending=context.pending or None,
                             incremental=context.incremental or None,
                             completed=context.completed or None,
                         )
-                    )
 
-                    yield subsequent_incremental_execution_result
+                        if not has_next:
+                            return
+                    batch = await incremental_graph.next_completed_batch()
         finally:
             await self._stop_async_iterators()