separate_operations: simplify algorithm

Replicates graphql/graphql-js@e2a6f23

Author: Cito

src/graphql/utilities/separate_operations.py

@@ -1,9 +1,8 @@
 from collections import defaultdict
-from typing import Dict, List, Set
+from typing import DefaultDict, Dict, List, Set
 
 from ..language import (
     DocumentNode,
-    ExecutableDefinitionNode,
     FragmentDefinitionNode,
     OperationDefinitionNode,
     Visitor,
@@ -13,7 +12,7 @@
 __all__ = ["separate_operations"]
 
 
-DepGraph = Dict[str, Set[str]]
+DepGraph = DefaultDict[str, Set[str]]
 
 
 def separate_operations(document_ast: DocumentNode) -> Dict[str, DocumentNode]:
@@ -23,13 +22,10 @@ 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.
     """
-
     # Populate metadata and build a dependency graph.
     visitor = SeparateOperations()
     visit(document_ast, visitor)
     operations = visitor.operations
-    fragments = visitor.fragments
-    positions = visitor.positions
     dep_graph = visitor.dep_graph
 
     # For each operation, produce a new synthesized AST which includes only what is
@@ -40,44 +36,42 @@ def separate_operations(document_ast: DocumentNode) -> Dict[str, DocumentNode]:
         dependencies: Set[str] = set()
         collect_transitive_dependencies(dependencies, dep_graph, operation_name)
 
-        # The list of definition nodes to be included for this operation, sorted to
-        # retain the same order as the original document.
-        definitions: List[ExecutableDefinitionNode] = [operation]
-        for name in dependencies:
-            definitions.append(fragments[name])
-            definitions.sort(key=lambda n: positions.get(n, 0))
-
-        separated_document_asts[operation_name] = DocumentNode(definitions=definitions)
+        # The list of definition nodes to be included for this operation, sorted
+        # to retain the same order as the original document.
+        separated_document_asts[operation_name] = DocumentNode(
+            definitions=[
+                node
+                for node in document_ast.definitions
+                if node is operation
+                or (
+                    isinstance(node, FragmentDefinitionNode)
+                    and node.name.value in dependencies
+                )
+            ]
+        )
 
     return separated_document_asts
 
 
-# noinspection PyAttributeOutsideInit
 class SeparateOperations(Visitor):
+    operations: List[OperationDefinitionNode]
+    dep_graph: DepGraph
+    from_name: str
+
     def __init__(self):
         super().__init__()
-        self.operations: List[OperationDefinitionNode] = []
-        self.fragments: Dict[str, FragmentDefinitionNode] = {}
-        self.positions: Dict[ExecutableDefinitionNode, int] = {}
-        self.dep_graph: DepGraph = defaultdict(set)
-        self.from_name: str
-        self.idx = 0
+        self.operations = []
+        self.dep_graph = defaultdict(set)
 
     def enter_operation_definition(self, node, *_args):
         self.from_name = op_name(node)
         self.operations.append(node)
-        self.positions[node] = self.idx
-        self.idx += 1
 
     def enter_fragment_definition(self, node, *_args):
         self.from_name = node.name.value
-        self.fragments[self.from_name] = node
-        self.positions[node] = self.idx
-        self.idx += 1
 
     def enter_fragment_spread(self, node, *_args):
-        to_name = node.name.value
-        self.dep_graph[self.from_name].add(to_name)
+        self.dep_graph[self.from_name].add(node.name.value)
 
 
 def op_name(operation: OperationDefinitionNode) -> str: