refactor(incremental): simplify incremental graph by allowing mutations

Replicates graphql/graphql-js@f6227a8

Author: Cito

src/graphql/execution/execute.py

@@ -2336,7 +2336,7 @@ def add_new_deferred_fragments(
 
         # Instantiate the new record.
         deferred_fragment_record = DeferredFragmentRecord(
-            parent, path, new_defer_usage.label
+            path, new_defer_usage.label, parent
         )
 
         # Update the map.

src/graphql/execution/incremental_graph.py

@@ -13,20 +13,19 @@
 )
 from collections import deque
 from contextlib import suppress
-from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, cast
 
 from ..pyutils import BoxedAwaitableOrValue, Undefined, is_awaitable
 from .types import (
     BareStreamItemsResult,
     StreamItemsResult,
     StreamRecord,
+    is_deferred_fragment_record,
     is_deferred_grouped_field_set_record,
 )
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Awaitable, Generator, Iterable, Sequence
-    from typing import TypeGuard
 
     from ..error.graphql_error import GraphQLError
     from .types import (
@@ -41,40 +40,13 @@
 __all__ = ["IncrementalGraph"]
 
 
-@dataclass(frozen=True, repr=False, eq=False)
-class DeferredFragmentNode:
-    """A node representing a deferred fragment in the incremental graph."""
-
-    deferred_fragment_record: DeferredFragmentRecord
-    deferred_grouped_field_set_records: dict[DeferredGroupedFieldSetRecord, None] = (
-        field(default_factory=dict)
-    )
-    reconcilable_results: dict[ReconcilableDeferredGroupedFieldSetResult, None] = field(
-        default_factory=dict
-    )
-    children: dict[SubsequentResultNode, None] = field(default_factory=dict)
-
-
-SubsequentResultNode = DeferredFragmentNode | StreamRecord
-
-
-def is_deferred_fragment_node(
-    node: SubsequentResultNode | None,
-) -> TypeGuard[DeferredFragmentNode]:
-    """Check whether the given node is a deferred fragment node."""
-    return isinstance(node, DeferredFragmentNode)
-
-
 class IncrementalGraph:
     """Helper class to execute incremental Graphs.
 
     For internal use only.
     """
 
-    _root_nodes: dict[SubsequentResultNode, None]
-    _deferred_fragment_nodes: dict[DeferredFragmentRecord, DeferredFragmentNode]
-    _new_pending: dict[SubsequentResultNode, None]
-    _new_incremental_data_records: dict[IncrementalDataRecord, None]
+    _root_nodes: dict[SubsequentResultRecord, None]
     _completed_queue: list[IncrementalDataRecordResult]
     _next_queue: list[Future[Iterable[IncrementalDataRecordResult]]]
 
@@ -83,9 +55,6 @@ class IncrementalGraph:
     def __init__(self) -> None:
         """Initialize the IncrementalGraph."""
         self._root_nodes = {}
-        self._deferred_fragment_nodes = {}
-        self._new_pending = {}
-        self._new_incremental_data_records = {}
         self._completed_queue = []
         self._next_queue = []
         self._tasks = set()
@@ -94,7 +63,7 @@ def get_new_root_nodes(
         self, incremental_data_records: Sequence[IncrementalDataRecord]
     ) -> list[SubsequentResultRecord]:
         """Get new root nodes."""
-        initial_result_children: dict[SubsequentResultNode, None] = {}
+        initial_result_children: dict[SubsequentResultRecord, None] = {}
         self._add_incremental_data_records(
             incremental_data_records, None, initial_result_children
         )
@@ -104,17 +73,19 @@ def add_completed_reconcilable_deferred_grouped_field_set(
         self, reconcilable_result: ReconcilableDeferredGroupedFieldSetResult
     ) -> None:
         """Add a completed reconcilable deferred grouped field set result."""
-        record = reconcilable_result.deferred_grouped_field_set_record
-        deferred = record.deferred_fragment_records
-        for defererred_fragment_node in self._fragments_to_nodes(deferred):
-            del defererred_fragment_node.deferred_grouped_field_set_records[
-                reconcilable_result.deferred_grouped_field_set_record
+        deferred_record = reconcilable_result.deferred_grouped_field_set_record
+        deferred_records = deferred_record.deferred_fragment_records
+        for defererred_fragment_record in deferred_records:
+            del defererred_fragment_record.deferred_grouped_field_set_records[
+                deferred_record
             ]
-            defererred_fragment_node.reconcilable_results[reconcilable_result] = None
+            defererred_fragment_record.reconcilable_results[reconcilable_result] = None
 
         incremental_data_records = reconcilable_result.incremental_data_records
         if incremental_data_records is not None:
-            self._add_incremental_data_records(incremental_data_records, deferred)
+            self._add_incremental_data_records(
+                incremental_data_records, deferred_records
+            )
 
     async def completed_incremental_data(
         self,
@@ -150,26 +121,22 @@ def complete_deferred_fragment(
         | None
     ):
         """Complete a deferred fragment."""
-        deferred_fragment_nodes = self._deferred_fragment_nodes
-        try:
-            deferred_fragment_node = deferred_fragment_nodes[deferred_fragment_record]
-        except KeyError:  # pragma: no cover
+        if deferred_fragment_record not in self._root_nodes:
+            return None  # pragma: no cover
+        if deferred_fragment_record.deferred_grouped_field_set_records:
             return None
-        if deferred_fragment_node.deferred_grouped_field_set_records:
-            return None
-        reconcilable_results = list(deferred_fragment_node.reconcilable_results)
-        self._remove_root_node(deferred_fragment_node)
+        reconcilable_results = list(deferred_fragment_record.reconcilable_results)
+        self._remove_root_node(deferred_fragment_record)
         for reconcilable_result in reconcilable_results:
-            record = reconcilable_result.deferred_grouped_field_set_record
-            for other_deferred_fragment_node in self._fragments_to_nodes(
-                record.deferred_fragment_records
-            ):
+            deferred_record = reconcilable_result.deferred_grouped_field_set_record
+            deferred_records = deferred_record.deferred_fragment_records
+            for other_deferred_fragment_record in deferred_records:
                 with suppress(KeyError):
-                    del other_deferred_fragment_node.reconcilable_results[
+                    del other_deferred_fragment_record.reconcilable_results[
                         reconcilable_result
                     ]
         new_root_nodes = self._promote_non_empty_to_root(
-            deferred_fragment_node.children
+            deferred_fragment_record.children
         )
         return new_root_nodes, reconcilable_results
 
@@ -178,16 +145,9 @@ def remove_deferred_fragment(
         deferred_fragment_record: DeferredFragmentRecord,
     ) -> bool:
         """Check if deferred fragment exists and remove it in that case."""
-        deferred_fragment_nodes = self._deferred_fragment_nodes
-        try:
-            deferred_fragment_node = deferred_fragment_nodes[deferred_fragment_record]
-        except KeyError:  # pragma: no cover
+        if deferred_fragment_record not in self._root_nodes:
             return False
-        self._remove_root_node(deferred_fragment_node)
-        del deferred_fragment_nodes[deferred_fragment_record]
-        for child in deferred_fragment_node.children:  # pragma: no cover
-            if is_deferred_fragment_node(child):
-                self.remove_deferred_fragment(child.deferred_fragment_record)
+        self._remove_root_node(deferred_fragment_record)
         return True
 
     def remove_stream(self, stream_record: StreamRecord) -> None:
@@ -199,28 +159,29 @@ def stop_incremental_data(self) -> None:
         for future in self._next_queue:
             future.cancel()  # pragma: no cover
 
-    def _remove_root_node(self, subsequent_result_node: SubsequentResultNode) -> None:
+    def _remove_root_node(
+        self, subsequent_result_record: SubsequentResultRecord
+    ) -> None:
         """Remove root node."""
-        del self._root_nodes[subsequent_result_node]
+        del self._root_nodes[subsequent_result_record]
         if not self._root_nodes:
             self.stop_incremental_data()
 
     def _add_incremental_data_records(
         self,
         incremental_data_records: Sequence[IncrementalDataRecord],
         parents: Sequence[DeferredFragmentRecord] | None = None,
-        initial_result_children: dict[SubsequentResultNode, None] | None = None,
+        initial_result_children: dict[SubsequentResultRecord, None] | None = None,
     ) -> None:
         """Add incremental data records."""
         for incremental_data_record in incremental_data_records:
             if is_deferred_grouped_field_set_record(incremental_data_record):
-                for (
-                    deferred_fragment_record
-                ) in incremental_data_record.deferred_fragment_records:
-                    deferred_fragment_node = self._add_deferred_fragment_node(
+                deferred_records = incremental_data_record.deferred_fragment_records
+                for deferred_fragment_record in deferred_records:
+                    self._add_deferred_fragment_node(
                         deferred_fragment_record, initial_result_children
                     )
-                    deferred_fragment_node.deferred_grouped_field_set_records[
+                    deferred_fragment_record.deferred_grouped_field_set_records[
                         incremental_data_record
                     ] = None
                 if self._completes_root_node(incremental_data_record):
@@ -234,23 +195,21 @@ def _add_incremental_data_records(
                 ] = None
             else:
                 for parent in parents:
-                    deferred_fragment_node = self._add_deferred_fragment_node(
-                        parent, initial_result_children
+                    self._add_deferred_fragment_node(parent, initial_result_children)
+                    parent.children[cast("StreamRecord", incremental_data_record)] = (
+                        None
                     )
-                    deferred_fragment_node.children[
-                        cast("StreamRecord", incremental_data_record)
-                    ] = None
 
     def _promote_non_empty_to_root(
-        self, maybe_empty_new_root_nodes: dict[SubsequentResultNode, None]
+        self, maybe_empty_new_root_nodes: dict[SubsequentResultRecord, None]
     ) -> list[SubsequentResultRecord]:
         """Promote non-empty nodes to root nodes."""
         new_root_nodes: list[SubsequentResultRecord] = []
         # use a deque to simulate how JavaScripts iterates over a changing set
         unprocessed_nodes = deque(maybe_empty_new_root_nodes)
         while unprocessed_nodes:
             node = unprocessed_nodes.popleft()
-            if is_deferred_fragment_node(node):
+            if is_deferred_fragment_record(node):
                 if node.deferred_grouped_field_set_records:
                     for (
                         deferred_grouped_field_set_record
@@ -262,9 +221,8 @@ def _promote_non_empty_to_root(
                                 deferred_grouped_field_set_record
                             )
                     self._root_nodes[node] = None
-                    new_root_nodes.append(node.deferred_fragment_record)
+                    new_root_nodes.append(node)
                     continue
-                del self._deferred_fragment_nodes[node.deferred_fragment_record]
                 for child in node.children:
                     if child not in maybe_empty_new_root_nodes:  # pragma: no branch
                         maybe_empty_new_root_nodes[cast("StreamRecord", child)] = None
@@ -280,54 +238,26 @@ def _completes_root_node(
     ) -> bool:
         """Check whether the given record completes a root node."""
         root_nodes = self._root_nodes
-        return any(
-            node in root_nodes
-            for node in self._fragments_to_nodes(
-                deferred_grouped_field_set_record.deferred_fragment_records
-            )
-        )
-
-    def _fragments_to_nodes(
-        self,
-        deferred_fragment_records: Sequence[DeferredFragmentRecord],
-    ) -> list[DeferredFragmentNode]:
-        """Get deferred fragment nodes for the given records."""
-        return [
-            node
-            for node in (
-                self._deferred_fragment_nodes.get(deferred_fragment_record)
-                for deferred_fragment_record in deferred_fragment_records
-            )
-            if is_deferred_fragment_node(node)
-        ]
+        deferred_records = deferred_grouped_field_set_record.deferred_fragment_records
+        return any(record in root_nodes for record in deferred_records)
 
     def _add_deferred_fragment_node(
         self,
         deferred_fragment_record: DeferredFragmentRecord,
-        initial_result_children: dict[SubsequentResultNode, None] | None = None,
-    ) -> DeferredFragmentNode:
+        initial_result_children: dict[SubsequentResultRecord, None] | None = None,
+    ) -> None:
         """Add a deferred fragment node."""
-        try:
-            deferred_fragment_node = self._deferred_fragment_nodes[
-                deferred_fragment_record
-            ]
-        except KeyError as key_error:
-            deferred_fragment_node = DeferredFragmentNode(deferred_fragment_record)
-            self._deferred_fragment_nodes[deferred_fragment_record] = (
-                deferred_fragment_node
-            )
-            parent = deferred_fragment_record.parent
-            if parent is None:
+        if deferred_fragment_record in self._root_nodes:
+            return
+        parent = deferred_fragment_record.parent
+        if parent is None:
+            if initial_result_children is None:  # pragma: no cover
                 msg = "Invalid state while adding deferred fragment node."
-                if initial_result_children is None:  # pragma: no cover
-                    raise RuntimeError(msg) from key_error
-                initial_result_children[deferred_fragment_node] = None
-            else:
-                parent_node = self._add_deferred_fragment_node(
-                    parent, initial_result_children
-                )
-                parent_node.children[deferred_fragment_node] = None
-        return deferred_fragment_node
+                raise RuntimeError(msg)
+            initial_result_children[deferred_fragment_record] = None
+            return
+        parent.children[deferred_fragment_record] = None
+        self._add_deferred_fragment_node(parent, initial_result_children)
 
     def _on_deferred_grouped_field_set(
         self, deferred_grouped_field_set_record: DeferredGroupedFieldSetRecord

src/graphql/execution/types.py

@@ -51,6 +51,7 @@
     "SubsequentIncrementalExecutionResult",
     "SubsequentResultRecord",
     "is_cancellable_stream_record",
+    "is_deferred_fragment_record",
     "is_deferred_grouped_field_set_record",
     "is_deferred_grouped_field_set_result",
     "is_non_reconcilable_deferred_grouped_field_set_result",
@@ -796,36 +797,57 @@ def __init__(
 class DeferredFragmentRecord:
     """Deferred fragment record"""
 
-    parent: DeferredFragmentRecord | None
     path: Path | None
     label: str | None
     id: str | None
+    parent: DeferredFragmentRecord | None
+    deferred_grouped_field_set_records: dict[DeferredGroupedFieldSetRecord, None]
+    reconcilable_results: dict[ReconcilableDeferredGroupedFieldSetResult, None]
+    children: dict[SubsequentResultRecord, None]
 
-    __slots__ = "id", "label", "parent", "path"
+    __slots__ = (
+        "children",
+        "deferred_grouped_field_set_records",
+        "id",
+        "label",
+        "parent",
+        "path",
+        "reconcilable_results",
+    )
 
     def __init__(
         self,
-        parent: DeferredFragmentRecord | None = None,
         path: Path | None = None,
         label: str | None = None,
+        parent: DeferredFragmentRecord | None = None,
     ) -> None:
-        self.parent = parent
         self.path = path
         self.label = label
+        self.parent = parent
         self.id = None
+        self.deferred_grouped_field_set_records = {}
+        self.reconcilable_results = {}
+        self.children = {}
 
     def __repr__(self) -> str:
         name = self.__class__.__name__
         args: list[str] = []
-        if self.parent:
-            args.append("parent")
         if self.path:
             args.append(f"path={self.path.as_list()!r}")
         if self.label:
             args.append(f"label={self.label!r}")
+        if self.parent:
+            args.append("parent")
         return f"{name}({', '.join(args)})"
 
 
+def is_deferred_fragment_record(
+    subsequent_result_record: SubsequentResultRecord,
+) -> TypeGuard[DeferredFragmentRecord]:
+    """Check if the subsequent result record is a deferred fragment record."""
+    return isinstance(subsequent_result_record, DeferredFragmentRecord)
+
+
 class StreamItemResult(NamedTuple):
     """Stream item result"""