refactor(incremental): introduce BoxedPromiseOrValue to save awaited results

Replicates graphql/graphql-js@062785e

Author: Cito

src/graphql/execution/execute.py

@@ -6,6 +6,7 @@
     CancelledError,
     TimeoutError,  # only needed for Python < 3.11  # noqa: A004
     ensure_future,
+    sleep,
 )
 from contextlib import suppress
 from copy import copy
@@ -40,6 +41,7 @@
 )
 from ..pyutils import (
     AwaitableOrValue,
+    BoxedAwaitableOrValue,
     Path,
     RefMap,
     Undefined,
@@ -1627,7 +1629,8 @@ def execute_deferred_grouped_field_sets(
             )
 
             deferred_grouped_field_set_record = DeferredGroupedFieldSetRecord(
-                deferred_fragment_records, cast("DeferredGroupedFieldSetResult", None)
+                deferred_fragment_records,
+                cast("BoxedAwaitableOrValue[DeferredGroupedFieldSetResult]", None),
             )
 
             if should_defer(parent_defer_usages, defer_usage_set):
@@ -1650,10 +1653,12 @@ async def executor(
                         return await result
                     return cast("DeferredGroupedFieldSetResult", result)
 
-                deferred_grouped_field_set_record.result = executor(
-                    deferred_grouped_field_set_record,
-                    grouped_field_set,
-                    defer_usage_set,
+                deferred_grouped_field_set_record.result = BoxedAwaitableOrValue(
+                    executor(
+                        deferred_grouped_field_set_record,
+                        grouped_field_set,
+                        defer_usage_set,
+                    )
                 )
             else:
                 executed = self.execute_deferred_grouped_field_set(
@@ -1665,7 +1670,9 @@ async def executor(
                     IncrementalContext(defer_usage_set),
                     defer_map,
                 )
-                deferred_grouped_field_set_record.result = executed
+                deferred_grouped_field_set_record.result = BoxedAwaitableOrValue(
+                    executed
+                )
 
             append_record(deferred_grouped_field_set_record)
 
@@ -1743,59 +1750,71 @@ async def await_result() -> StreamItemsResult:
             path = stream_record.path
             initial_path = Path(path, initial_index, None)
 
-            result = self.complete_stream_items(
-                stream_record,
-                initial_path,
-                initial_item,
-                IncrementalContext(),
-                field_group,
-                info,
-                item_type,
+            result: BoxedAwaitableOrValue[StreamItemsResult] = BoxedAwaitableOrValue(
+                self.complete_stream_items(
+                    stream_record,
+                    initial_path,
+                    initial_item,
+                    IncrementalContext(),
+                    field_group,
+                    info,
+                    item_type,
+                )
             )
             first_stream_items = StreamItemsRecord(stream_record, result)
             current_stream_items = first_stream_items
             current_index = initial_index
             errored_synchronously = False
             for item in iterator:
-                if not is_awaitable(result) and not is_reconcilable_stream_items_result(
-                    result  # type: ignore
+                value = result.value
+                if not is_awaitable(value) and not is_reconcilable_stream_items_result(
+                    value
                 ):
                     errored_synchronously = True
                     break
                 current_index += 1
                 current_path = Path(path, current_index, None)
-                result = self.complete_stream_items(
-                    stream_record,
-                    current_path,
-                    item,
-                    IncrementalContext(),
-                    field_group,
-                    info,
-                    item_type,
+                result = BoxedAwaitableOrValue(
+                    self.complete_stream_items(
+                        stream_record,
+                        current_path,
+                        item,
+                        IncrementalContext(),
+                        field_group,
+                        info,
+                        item_type,
+                    )
                 )
 
                 next_stream_items = StreamItemsRecord(stream_record, result)
-                current_stream_items.result = prepend_next_stream_items(
-                    current_stream_items.result, next_stream_items
+                current_stream_items.result = BoxedAwaitableOrValue(
+                    prepend_next_stream_items(
+                        current_stream_items.result.value, next_stream_items
+                    )
                 )
                 current_stream_items = next_stream_items
 
             # If a non-reconcilable stream items result was encountered,
             # then the stream terminates in error. Otherwise, add a stream terminator.
             if not errored_synchronously:
-                current_stream_items.result = prepend_next_stream_items(
-                    current_stream_items.result,
-                    StreamItemsRecord(
-                        stream_record, TerminatingStreamItemsResult(stream_record)
-                    ),
+                current_stream_items.result = BoxedAwaitableOrValue(
+                    prepend_next_stream_items(
+                        current_stream_items.result.value,
+                        StreamItemsRecord(
+                            stream_record,
+                            BoxedAwaitableOrValue(
+                                TerminatingStreamItemsResult(stream_record),
+                            ),
+                        ),
+                    )
                 )
 
-            result = first_stream_items.result
-            if is_awaitable(result):
-                return await result
-            return cast("StreamItemsResult", result)
+            value = first_stream_items.result.value
+            if is_awaitable(value):
+                return await value
+            return value
 
-        return StreamItemsRecord(stream_record, await_result())
+        return StreamItemsRecord(stream_record, BoxedAwaitableOrValue(await_result()))
 
     def first_async_stream_items(
         self,
@@ -1810,14 +1829,16 @@ def first_async_stream_items(
         """Get the first async stream items."""
         return StreamItemsRecord(
             stream_record,
-            self.get_next_async_stream_items_result(
-                stream_record,
-                path,
-                initial_index,
-                async_iterator,
-                field_group,
-                info,
-                item_type,
+            BoxedAwaitableOrValue(
+                self.get_next_async_stream_items_result(
+                    stream_record,
+                    path,
+                    initial_index,
+                    async_iterator,
+                    field_group,
+                    info,
+                    item_type,
+                )
             ),
         )
 
@@ -1856,14 +1877,23 @@ async def get_next_async_stream_items_result(
 
         next_stream_items_record = StreamItemsRecord(
             stream_record,
-            self.get_next_async_stream_items_result(
-                stream_record, path, index, async_iterator, field_group, info, item_type
+            BoxedAwaitableOrValue(
+                self.get_next_async_stream_items_result(
+                    stream_record,
+                    path,
+                    index,
+                    async_iterator,
+                    field_group,
+                    info,
+                    item_type,
+                )
             ),
         )
 
         result = self.prepend_next_stream_items(result, next_stream_items_record)
 
         if self.is_awaitable(result):
+            await sleep(0)
             return await result
         return cast("StreamItemsResult", result)
 

src/graphql/execution/incremental_graph.py

@@ -2,7 +2,14 @@
 
 from __future__ import annotations
 
-from asyncio import CancelledError, Future, Task, ensure_future
+from asyncio import (
+    CancelledError,
+    Future,
+    Task,
+    ensure_future,
+    get_running_loop,
+    isfuture,
+)
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -20,8 +27,6 @@
     is_deferred_grouped_field_set_record,
 )
 
-from ..pyutils.is_awaitable import is_awaitable
-
 if TYPE_CHECKING:
     from graphql.execution.types import (
         DeferredFragmentRecord,
@@ -163,17 +168,17 @@ def get_new_pending(self) -> list[SubsequentResultRecord]:
 
         enqueue = self._enqueue
         for incremental_data_record in _new_incremental_data_records:
-            result = incremental_data_record.result
-            if is_awaitable(result):
+            value = incremental_data_record.result.value
+            if isfuture(value):
 
-                async def enqueue_incremental(
-                    result: Awaitable[IncrementalDataRecordResult],
+                async def enqueue_later(
+                    value: Awaitable[IncrementalDataRecordResult],
                 ) -> None:
-                    enqueue(await result)
+                    enqueue(await value)
 
-                self._add_task(enqueue_incremental(result))
+                self._add_task(enqueue_later(value))
             else:
-                enqueue(result)  # type: ignore
+                enqueue(value)
         _new_incremental_data_records.clear()
 
         return new_pending
@@ -182,12 +187,15 @@ async def completed_incremental_data(
         self,
     ) -> AsyncGenerator[Iterable[IncrementalDataRecordResult], None]:
         """Asynchronously yield completed incremental data record results."""
+        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]] = Future()
+                future: Future[Iterable[IncrementalDataRecordResult]] = (
+                    loop.create_future()
+                )
                 self._next_queue.append(future)
                 try:
                     yield await future

src/graphql/execution/incremental_publisher.py

@@ -156,7 +156,7 @@ async def _subscribe(
                 for completed_result in completed_results:
                     await handle_completed_incremental_data(completed_result, context)
 
-                if context.incremental or context.completed:
+                if context.incremental or context.completed:  # pragma: no branch
                     has_next = check_has_next()
 
                     if not has_next:
@@ -192,6 +192,7 @@ async def _handle_completed_incremental_data(
         completed_incremental_data: IncrementalDataRecordResult,
         context: SubsequentIncrementalExecutionResultContext,
     ) -> None:
+        """Handle completed incremental data."""
         if is_deferred_grouped_field_set_result(completed_incremental_data):
             self._handle_completed_deferred_grouped_field_set(
                 completed_incremental_data, context
@@ -203,7 +204,6 @@ async def _handle_completed_incremental_data(
             await self._handle_completed_stream_items(
                 completed_incremental_data, context
             )
-
         new_pending = self._incremental_graph.get_new_pending()
         context.pending.extend(self._pending_sources_to_results(new_pending))
 

src/graphql/execution/types.py

@@ -19,7 +19,7 @@
 
 if TYPE_CHECKING:
     from ..error import GraphQLError, GraphQLFormattedError
-    from ..pyutils import AwaitableOrValue, Path
+    from ..pyutils import BoxedAwaitableOrValue, Path
 
     try:
         from typing import TypeGuard
@@ -783,14 +783,14 @@ class DeferredGroupedFieldSetRecord:
     """Deferred grouped field set record"""
 
     deferred_fragment_records: list[DeferredFragmentRecord]
-    result: AwaitableOrValue[DeferredGroupedFieldSetResult]
+    result: BoxedAwaitableOrValue[DeferredGroupedFieldSetResult]
 
     __slots__ = "deferred_fragment_records", "result"
 
     def __init__(
         self,
         deferred_fragment_records: list[DeferredFragmentRecord],
-        result: AwaitableOrValue[DeferredGroupedFieldSetResult],
+        result: BoxedAwaitableOrValue[DeferredGroupedFieldSetResult],
     ) -> None:
         self.result = result
         self.deferred_fragment_records = deferred_fragment_records
@@ -923,12 +923,12 @@ class StreamItemsRecord:
     __slots__ = "result", "stream_record"
 
     stream_record: StreamRecord
-    result: AwaitableOrValue[StreamItemsResult]
+    result: BoxedAwaitableOrValue[StreamItemsResult]
 
     def __init__(
         self,
         stream_record: StreamRecord,
-        result: AwaitableOrValue[StreamItemsResult],
+        result: BoxedAwaitableOrValue[StreamItemsResult],
     ) -> None:
         self.stream_record = stream_record
         self.result = result

src/graphql/pyutils/__init__.py

@@ -9,6 +9,7 @@
 """
 
 from .async_reduce import async_reduce
+from .boxed_awaitable_or_value import BoxedAwaitableOrValue
 from .gather_with_cancel import gather_with_cancel
 from .convert_case import camel_to_snake, snake_to_camel
 from .cached_property import cached_property
@@ -39,6 +40,7 @@
 
 __all__ = [
     "AwaitableOrValue",
+    "BoxedAwaitableOrValue",
     "Description",
     "FrozenError",
     "Path",