separate_operations: distinguish query and fragment names

Replicates graphql/graphql-js@cd273ad

Author: Cito

src/graphql/utilities/separate_operations.py

@@ -1,19 +1,19 @@
-from collections import defaultdict
-from typing import Any, DefaultDict, Dict, List, Set
+from typing import Any, Dict, List, Set
 
 from ..language import (
     DocumentNode,
     FragmentDefinitionNode,
     FragmentSpreadNode,
     OperationDefinitionNode,
+    SelectionSetNode,
     Visitor,
     visit,
 )
 
 __all__ = ["separate_operations"]
 
 
-DepGraph = DefaultDict[str, Set[str]]
+DepGraph = Dict[str, List[str]]
 
 
 def separate_operations(document_ast: DocumentNode) -> Dict[str, DocumentNode]:
@@ -23,19 +23,31 @@ def separate_operations(document_ast: DocumentNode) -> Dict[str, DocumentNode]:
     fragments and returns a collection of AST documents each of which contains a single
     operation as well the fragment definitions it refers to.
     """
+    operations: List[OperationDefinitionNode] = []
+    dep_graph: DepGraph = {}
+
     # Populate metadata and build a dependency graph.
-    visitor = SeparateOperations()
-    visit(document_ast, visitor)
-    operations = visitor.operations
-    dep_graph = visitor.dep_graph
+    for definition_node in document_ast.definitions:
+        if isinstance(definition_node, OperationDefinitionNode):
+            operations.append(definition_node)
+        elif isinstance(
+            definition_node, FragmentDefinitionNode
+        ):  # pragma: no cover else
+            dep_graph[definition_node.name.value] = collect_dependencies(
+                definition_node.selection_set
+            )
 
     # For each operation, produce a new synthesized AST which includes only what is
     # necessary for completing that operation.
-    separated_document_asts = {}
+    separated_document_asts: Dict[str, DocumentNode] = {}
     for operation in operations:
-        operation_name = op_name(operation)
         dependencies: Set[str] = set()
-        collect_transitive_dependencies(dependencies, dep_graph, operation_name)
+
+        for fragment_name in collect_dependencies(operation.selection_set):
+            collect_transitive_dependencies(dependencies, dep_graph, fragment_name)
+
+        # Provides the empty string for anonymous operations.
+        operation_name = operation.name.value if operation.name else ""
 
         # The list of definition nodes to be included for this operation, sorted
         # to retain the same order as the original document.
@@ -54,36 +66,6 @@ def separate_operations(document_ast: DocumentNode) -> Dict[str, DocumentNode]:
     return separated_document_asts
 
 
-class SeparateOperations(Visitor):
-    operations: List[OperationDefinitionNode]
-    dep_graph: DepGraph
-    from_name: str
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.operations = []
-        self.dep_graph = defaultdict(set)
-
-    def enter_operation_definition(
-        self, node: OperationDefinitionNode, *_args: Any
-    ) -> None:
-        self.from_name = op_name(node)
-        self.operations.append(node)
-
-    def enter_fragment_definition(
-        self, node: FragmentDefinitionNode, *_args: Any
-    ) -> None:
-        self.from_name = node.name.value
-
-    def enter_fragment_spread(self, node: FragmentSpreadNode, *_args: Any) -> None:
-        self.dep_graph[self.from_name].add(node.name.value)
-
-
-def op_name(operation: OperationDefinitionNode) -> str:
-    """Provide the empty string for anonymous operations."""
-    return operation.name.value if operation.name else ""
-
-
 def collect_transitive_dependencies(
     collected: Set[str], dep_graph: DepGraph, from_name: str
 ) -> None:
@@ -92,8 +74,28 @@ def collect_transitive_dependencies(
     From a dependency graph, collects a list of transitive dependencies by recursing
     through a dependency graph.
     """
-    immediate_deps = dep_graph[from_name]
-    for to_name in immediate_deps:
-        if to_name not in collected:
-            collected.add(to_name)
-            collect_transitive_dependencies(collected, dep_graph, to_name)
+    if from_name not in collected:
+        collected.add(from_name)
+
+        immediate_deps = dep_graph.get(from_name)
+        if immediate_deps is not None:
+            for to_name in immediate_deps:
+                collect_transitive_dependencies(collected, dep_graph, to_name)
+
+
+class DependencyCollector(Visitor):
+    dependencies: List[str]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.dependencies = []
+        self.add_dependency = self.dependencies.append
+
+    def enter_fragment_spread(self, node: FragmentSpreadNode, *_args: Any) -> None:
+        self.add_dependency(node.name.value)
+
+
+def collect_dependencies(selection_set: SelectionSetNode) -> List[str]:
+    collector = DependencyCollector()
+    visit(selection_set, collector)
+    return collector.dependencies