Bring visitor_keys back

Replicates graphql/graphql-js@0091421

Author: Cito

src/graphql/__init__.py

@@ -187,6 +187,7 @@
     ParallelVisitor,
     Visitor,
     VisitorAction,
+    VisitorKeyMap,
     BREAK,
     SKIP,
     REMOVE,
@@ -553,6 +554,7 @@
     "TypeInfoVisitor",
     "Visitor",
     "VisitorAction",
+    "VisitorKeyMap",
     "BREAK",
     "SKIP",
     "REMOVE",

src/graphql/language/__init__.py

@@ -23,6 +23,7 @@
     Visitor,
     ParallelVisitor,
     VisitorAction,
+    VisitorKeyMap,
     BREAK,
     SKIP,
     REMOVE,
@@ -126,6 +127,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": (
@@ -76,11 +79,6 @@ class VisitorActionEnum(Enum):
         "directives",
         "selection_set",
     ),
-    "int_value": (),
-    "float_value": (),
-    "string_value": (),
-    "boolean_value": (),
-    "enum_value": (),
     "list_value": ("values",),
     "object_value": ("fields",),
     "object_field": ("name", "value"),
@@ -212,7 +210,9 @@ class Stack(NamedTuple):
     prev: Any  # 'Stack' (python/mypy/issues/731)
 
 
-def visit(root: Node, visitor: Visitor) -> Any:
+def visit(
+    root: Node, visitor: Visitor, visitor_keys: Optional[VisitorKeyMap] = None
+) -> Any:
     """Visit each node in an AST.
 
     :func:`~.visit` will walk through an AST using a depth-first traversal, calling the
@@ -227,11 +227,16 @@ def visit(root: Node, visitor: Visitor) -> Any:
     When using :func:`~.visit` to edit an AST, the original AST will not be modified,
     and a new version of the AST with the changes applied will be returned from the
     visit function.
+
+    To customize the node attributes to be used for traversal, you can provide a
+    dictionary visitor_keys mapping node kinds to node attributes.
     """
     if not isinstance(root, Node):
         raise TypeError(f"Not an AST Node: {inspect(root)}.")
     if not isinstance(visitor, Visitor):
         raise TypeError(f"Not an AST Visitor: {inspect(visitor)}.")
+    if visitor_keys is None:
+        visitor_keys = QUERY_DOCUMENT_KEYS
     stack: Any = None
     in_array = isinstance(root, list)
     keys: Tuple[Node, ...] = (root,)
@@ -326,7 +331,7 @@ def visit(root: Node, visitor: Visitor) -> Any:
         else:
             stack = Stack(in_array, idx, keys, edits, stack)
             in_array = isinstance(node, list)
-            keys = node if in_array else QUERY_DOCUMENT_KEYS.get(node.kind, ())
+            keys = node if in_array else visitor_keys.get(node.kind, ())
             idx = -1
             edits = []
             if parent:

tests/language/test_visitor.py

@@ -17,6 +17,7 @@
     SKIP,
     ParallelVisitor,
     Visitor,
+    VisitorKeyMap,
 )
 from graphql.pyutils import FrozenList
 
@@ -588,6 +589,41 @@ def leave(node, *_args):
             ["leave", "document", None],
         ]
 
+    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):
+                visited.append(["enter", node.kind, get_value(node)])
+
+            @staticmethod
+            def leave(node, *_args):
+                visited.append(["leave", node.kind, get_value(node)])
+
+        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", "name", "ExampleOperation"],
+            ["leave", "name", "ExampleOperation"],
+            ["leave", "operation_definition", None],
+            ["leave", "document", None],
+        ]
+
     def cannot_define_visitor_with_unknown_ast_nodes():
         with raises(TypeError) as exc_info: