Backport VisitorKeyMap type

Replicates graphql/graphql-js@201ed90

Author: Cito

src/graphql/__init__.py

@@ -182,6 +182,7 @@
     ParallelVisitor,
     Visitor,
     VisitorAction,
+    VisitorKeyMap,
     BREAK,
     SKIP,
     REMOVE,
@@ -540,6 +541,7 @@
     "TypeInfoVisitor",
     "Visitor",
     "VisitorAction",
+    "VisitorKeyMap",
     "BREAK",
     "SKIP",
     "REMOVE",

src/graphql/language/__init__.py

@@ -23,6 +23,7 @@
     Visitor,
     ParallelVisitor,
     VisitorAction,
+    VisitorKeyMap,
     BREAK,
     SKIP,
     REMOVE,
@@ -117,6 +118,7 @@
     "Visitor",
     "ParallelVisitor",
     "VisitorAction",
+    "VisitorKeyMap",
     "BREAK",
     "SKIP",
     "REMOVE",

src/graphql/language/visitor.py

@@ -21,6 +21,7 @@
     "Visitor",
     "ParallelVisitor",
     "VisitorAction",
+    "VisitorKeyMap",
     "visit",
     "BREAK",
     "SKIP",
@@ -51,8 +52,10 @@ class VisitorActionEnum(Enum):
 REMOVE = VisitorActionEnum.REMOVE
 IDLE = None
 
+VisitorKeyMap = Dict[str, Tuple[str, ...]]
+
 # Default map from visitor kinds to their traversable node attributes:
-QUERY_DOCUMENT_KEYS: Dict[str, Tuple[str, ...]] = {
+QUERY_DOCUMENT_KEYS: VisitorKeyMap = {
     "name": (),
     "document": ("definitions",),
     "operation_definition": (
@@ -168,6 +171,9 @@ def leave(self, node, key, parent, path, ancestors):
     nodes, you would defined the methods ``enter_field()`` and/or ``leave_field()``,
     with the same signature as above. If no kind specific method has been defined
     for a given node, the generic method is called.
+
+    To customize the node attributes to be used for traversal, you can provide a
+    dictionary visitor_keys mapping node kinds to node attributes.
     """
 
     # Provide special return values as attributes
@@ -215,9 +221,7 @@ class Stack(NamedTuple):
 
 
 def visit(
-    root: Node,
-    visitor: Visitor,
-    visitor_keys: Optional[Dict[str, Tuple[str, ...]]] = None,
+    root: Node, visitor: Visitor, visitor_keys: Optional[VisitorKeyMap] = None
 ) -> Any:
     """Visit each node in an AST.
 

tests/language/test_visitor.py

@@ -1,6 +1,6 @@
 from copy import copy
 from functools import partial
-from typing import cast, Dict, List, Optional, Tuple
+from typing import cast, List, Optional
 
 from pytest import mark, raises
 
@@ -17,8 +17,8 @@
     SKIP,
     ParallelVisitor,
     Visitor,
+    VisitorKeyMap,
 )
-from graphql.language.visitor import QUERY_DOCUMENT_KEYS
 from graphql.pyutils import FrozenList
 
 from ..fixtures import kitchen_sink_query  # noqa: F401
@@ -1019,13 +1019,14 @@ def leave(node, *_args):
             ["leave", "document", None],
         ]
 
-    def does_traverse_unknown_node_with_visitor_keys():
-        custom_query_document_keys: Dict[str, Tuple[str, ...]] = {
-            **QUERY_DOCUMENT_KEYS,
-            "custom_field": ("name", "selection_set"),
-        }
+    def visits_only_the_specified_kind_in_visitor_key_map():
         visited = []
 
+        visitor_key_map: VisitorKeyMap = {
+            "document": ("definitions",),
+            "operation_definition": ("name",),
+        }
+
         class TestVisitor(Visitor):
             @staticmethod
             def enter(node, *_args):
@@ -1035,26 +1036,20 @@ def enter(node, *_args):
             def leave(node, *_args):
                 visited.append(["leave", node.kind, get_value(node)])
 
-        visit(custom_ast, TestVisitor(), custom_query_document_keys)
+        example_document_ast = parse(
+            """
+            query ExampleOperation {
+              someField
+            }
+            """
+        )
+
+        visit(example_document_ast, TestVisitor(), visitor_key_map)
         assert visited == [
             ["enter", "document", None],
             ["enter", "operation_definition", None],
-            ["enter", "selection_set", None],
-            ["enter", "field", None],
-            ["enter", "name", "a"],
-            ["leave", "name", "a"],
-            ["leave", "field", None],
-            ["enter", "custom_field", None],
-            ["enter", "name", "b"],
-            ["leave", "name", "b"],
-            ["enter", "selection_set", None],
-            ["enter", "custom_field", None],
-            ["enter", "name", "c"],
-            ["leave", "name", "c"],
-            ["leave", "custom_field", None],
-            ["leave", "selection_set", None],
-            ["leave", "custom_field", None],
-            ["leave", "selection_set", None],
+            ["enter", "name", "ExampleOperation"],
+            ["leave", "name", "ExampleOperation"],
             ["leave", "operation_definition", None],
             ["leave", "document", None],
         ]