visitor: Speed up ParallelVisitor

Contrary to GraphQL.js, we still support custom node types.
However, we use a similar caching mechanism for the visit functions
to speed up the visit process, particularly for the ParallelVisitor.

Replicates graphql/graphql-js@f4efee9

Author: Cito

docs/conf.py

@@ -133,9 +133,11 @@
 enum.Enum
 traceback
 types.TracebackType
+EnterLeaveVisitor
 FormattedSourceLocation
 asyncio.events.AbstractEventLoop
 graphql.language.lexer.EscapeSequence
+graphql.language.visitor.EnterLeaveVisitor
 graphql.subscription.map_async_iterator.MapAsyncIterator
 graphql.type.schema.InterfaceImplementations
 graphql.validation.validation_context.VariableUsage

src/graphql/language/visitor.py

@@ -126,6 +126,13 @@ class VisitorActionEnum(Enum):
 }
 
 
+class EnterLeaveVisitor(NamedTuple):
+    """Visitor with functions for entering and leaving."""
+
+    enter: Optional[Callable[..., Optional[VisitorAction]]]
+    leave: Optional[Callable[..., Optional[VisitorAction]]]
+
+
 class Visitor:
     """Visitor that walks through an AST.
 
@@ -170,6 +177,8 @@ def leave(self, node, key, parent, path, ancestors):
     # Provide special return values as attributes
     BREAK, SKIP, REMOVE, IDLE = BREAK, SKIP, REMOVE, IDLE
 
+    enter_leave_map: Dict[str, EnterLeaveVisitor]
+
     def __init_subclass__(cls) -> None:
         """Verify that all defined handlers are valid."""
         super().__init_subclass__()
@@ -191,13 +200,34 @@ def __init_subclass__(cls) -> None:
                 ):
                     raise TypeError(f"Invalid AST node kind: {kind}.")
 
-    def get_visit_fn(self, kind: str, is_leaving: bool = False) -> Callable:
-        """Get the visit function for the given node kind and direction."""
-        method = "leave" if is_leaving else "enter"
-        visit_fn = getattr(self, f"{method}_{kind}", None)
-        if not visit_fn:
-            visit_fn = getattr(self, method, None)
-        return visit_fn
+    def __init__(self) -> None:
+        self.enter_leave_map = {}
+
+    def get_enter_leave_for_kind(self, kind: str) -> EnterLeaveVisitor:
+        """Given a node kind, return the EnterLeaveVisitor for that kind."""
+        try:
+            return self.enter_leave_map[kind]
+        except KeyError:
+            enter_fn = getattr(self, f"enter_{kind}", None)
+            if not enter_fn:
+                enter_fn = getattr(self, "enter", None)
+            leave_fn = getattr(self, f"leave_{kind}", None)
+            if not leave_fn:
+                leave_fn = getattr(self, "leave", None)
+            enter_leave = EnterLeaveVisitor(enter_fn, leave_fn)
+            self.enter_leave_map[kind] = enter_leave
+            return enter_leave
+
+    def get_visit_fn(
+        self, kind: str, is_leaving: bool = False
+    ) -> Optional[Callable[..., Optional[VisitorAction]]]:
+        """Get the visit function for the given node kind and direction.
+
+        .. deprecated:: 3.2
+           Please use ``get_enter_leave_for_kind`` instead. Will be removed in v3.3.
+        """
+        enter_leave = self.get_enter_leave_for_kind(kind)
+        return enter_leave.leave if is_leaving else enter_leave.enter
 
 
 class Stack(NamedTuple):
@@ -237,6 +267,7 @@ def visit(
         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,)
@@ -299,7 +330,8 @@ def visit(
         else:
             if not isinstance(node, Node):
                 raise TypeError(f"Invalid AST Node: {inspect(node)}.")
-            visit_fn = visitor.get_visit_fn(node.kind, is_leaving)
+            enter_leave = visitor.get_enter_leave_for_kind(node.kind)
+            visit_fn = enter_leave.leave if is_leaving else enter_leave.enter
             if visit_fn:
                 result = visit_fn(node, key, parent, path, ancestors)
 
@@ -357,39 +389,63 @@ class ParallelVisitor(Visitor):
 
     def __init__(self, visitors: Collection[Visitor]):
         """Create a new visitor from the given list of parallel visitors."""
+        super().__init__()
         self.visitors = visitors
         self.skipping: List[Any] = [None] * len(visitors)
 
-    def enter(self, node: Node, *args: Any) -> Optional[VisitorAction]:
-        skipping = self.skipping
-        for i, visitor in enumerate(self.visitors):
-            if not skipping[i]:
-                fn = visitor.get_visit_fn(node.kind)
-                if fn:
-                    result = fn(node, *args)
-                    if result is SKIP or result is False:
-                        skipping[i] = node
-                    elif result is BREAK or result is True:
-                        skipping[i] = BREAK
-                    elif result is not None:
-                        return result
-        return None
-
-    def leave(self, node: Node, *args: Any) -> Optional[VisitorAction]:
-        skipping = self.skipping
-        for i, visitor in enumerate(self.visitors):
-            if not skipping[i]:
-                fn = visitor.get_visit_fn(node.kind, is_leaving=True)
-                if fn:
-                    result = fn(node, *args)
-                    if result is BREAK or result is True:
-                        skipping[i] = BREAK
-                    elif (
-                        result is not None
-                        and result is not SKIP
-                        and result is not False
-                    ):
-                        return result
-            elif skipping[i] is node:
-                skipping[i] = None
-        return None
+    def get_enter_leave_for_kind(self, kind: str) -> EnterLeaveVisitor:
+        """Given a node kind, return the EnterLeaveVisitor for that kind."""
+        try:
+            return self.enter_leave_map[kind]
+        except KeyError:
+            has_visitor = False
+            enter_list: List[Optional[Callable[..., Optional[VisitorAction]]]] = []
+            leave_list: List[Optional[Callable[..., Optional[VisitorAction]]]] = []
+            for visitor in self.visitors:
+                enter, leave = visitor.get_enter_leave_for_kind(kind)
+                if not has_visitor and (enter or leave):
+                    has_visitor = True
+                enter_list.append(enter)
+                leave_list.append(leave)
+
+            if has_visitor:
+
+                def enter(node: Node, *args: Any) -> Optional[VisitorAction]:
+                    skipping = self.skipping
+                    for i, fn in enumerate(enter_list):
+                        if not skipping[i]:
+                            if fn:
+                                result = fn(node, *args)
+                                if result is SKIP or result is False:
+                                    skipping[i] = node
+                                elif result is BREAK or result is True:
+                                    skipping[i] = BREAK
+                                elif result is not None:
+                                    return result
+                    return None
+
+                def leave(node: Node, *args: Any) -> Optional[VisitorAction]:
+                    skipping = self.skipping
+                    for i, fn in enumerate(leave_list):
+                        if not skipping[i]:
+                            if fn:
+                                result = fn(node, *args)
+                                if result is BREAK or result is True:
+                                    skipping[i] = BREAK
+                                elif (
+                                    result is not None
+                                    and result is not SKIP
+                                    and result is not False
+                                ):
+                                    return result
+                        elif skipping[i] is node:
+                            skipping[i] = None
+                    return None
+
+            else:
+
+                enter = leave = None
+
+            enter_leave = EnterLeaveVisitor(enter, leave)
+            self.enter_leave_map[kind] = enter_leave
+            return enter_leave

src/graphql/utilities/type_info.py

@@ -299,12 +299,13 @@ class TypeInfoVisitor(Visitor):
     """A visitor which maintains a provided TypeInfo."""
 
     def __init__(self, type_info: "TypeInfo", visitor: Visitor):
+        super().__init__()
         self.type_info = type_info
         self.visitor = visitor
 
     def enter(self, node: Node, *args: Any) -> Any:
         self.type_info.enter(node)
-        fn = self.visitor.get_visit_fn(node.kind)
+        fn = self.visitor.get_enter_leave_for_kind(node.kind).enter
         if fn:
             result = fn(node, *args)
             if result is not None:
@@ -314,7 +315,7 @@ def enter(self, node: Node, *args: Any) -> Any:
             return result
 
     def leave(self, node: Node, *args: Any) -> Any:
-        fn = self.visitor.get_visit_fn(node.kind, is_leaving=True)
+        fn = self.visitor.get_enter_leave_for_kind(node.kind).leave
         result = fn(node, *args) if fn else None
         self.type_info.leave(node)
         return result

src/graphql/validation/rules/__init__.py

@@ -17,6 +17,7 @@ class ASTValidationRule(Visitor):
     context: ASTValidationContext
 
     def __init__(self, context: ASTValidationContext):
+        super().__init__()
         self.context = context
 
     def report_error(self, error: GraphQLError) -> None:

src/graphql/validation/validation_context.py

@@ -47,6 +47,7 @@ class VariableUsageVisitor(Visitor):
     usages: List[VariableUsage]
 
     def __init__(self, type_info: TypeInfo):
+        super().__init__()
         self.usages = []
         self._append_usage = self.usages.append
         self._type_info = type_info

tests/benchmarks/__init__.py

@@ -3,7 +3,7 @@
 Benchmarks are disabled (only executed as tests) by default in setup.cfg.
 You can enable them with --benchmark-enable if your want to execute them.
 
-E.g. in order to execute all the benchmarks with tox using Python 3.7::
+E.g. in order to execute all the benchmarks with tox using Python 3.9::
 
-    tox -e py37 -- -k benchmarks --benchmark-enable
+    tox -e py39 -- -k benchmarks --benchmark-enable
 """

tests/benchmarks/test_visit.py

@@ -0,0 +1,27 @@
+from graphql import parse
+from graphql.language import visit, Visitor, ParallelVisitor
+
+from ..fixtures import big_schema_sdl  # noqa: F401
+
+
+class TestVisitor(Visitor):
+    @staticmethod
+    def enter(*args):
+        pass
+
+    @staticmethod
+    def leave(*args):
+        pass
+
+
+def test_visit_all_ast_nodes(benchmark, big_schema_sdl):  # noqa: F811
+    document_ast = parse(big_schema_sdl)
+    visitor = TestVisitor()
+    benchmark(lambda: visit(document_ast, visitor))
+
+
+def test_visit_all_ast_nodes_in_parallel(benchmark, big_schema_sdl):  # noqa: F811
+    document_ast = parse(big_schema_sdl)
+    visitor = TestVisitor()
+    parallel_visitor = ParallelVisitor([visitor] * 50)
+    benchmark(lambda: visit(document_ast, parallel_visitor))

tests/language/test_visitor.py

@@ -214,6 +214,41 @@ def leave_field(node, *args):
                 "leave:document",
             ]
 
+    def has_get_enter_leave_for_kind_method():
+        class TestVisitor(Visitor):
+            @staticmethod
+            def enter(*args):
+                pass
+
+            @staticmethod
+            def enter_document(*args):
+                pass
+
+            @staticmethod
+            def leave(*args):
+                pass
+
+            @staticmethod
+            def leave_document(*args):
+                pass
+
+        visitor = TestVisitor()
+
+        assert visitor.get_enter_leave_for_kind("document") == (
+            visitor.enter_document,
+            visitor.leave_document,
+        )
+        assert visitor.get_enter_leave_for_kind("field") == (
+            visitor.enter,
+            visitor.leave,
+        )
+
+        # also test deprecated method
+        assert visitor.get_visit_fn("document") == visitor.enter_document
+        assert visitor.get_visit_fn("field") == visitor.enter
+        assert visitor.get_visit_fn("document", True) == visitor.leave_document
+        assert visitor.get_visit_fn("field", True) == visitor.leave
+
     def validates_path_argument():
         ast = parse("{ a }", no_location=True)
         visited = []
@@ -540,7 +575,10 @@ def leave_selection_set(*args):
             ["leave", "selection_set", None],
         ]
 
-    def visit_nodes_with_unknown_kinds_but_does_not_traverse_deeper():
+    def visit_nodes_with_custom_kinds_but_does_not_traverse_deeper():
+        # GraphQL.js removed support for unknown node types,
+        # but it is easy for us to add and support custom node types,
+        # so we keep allowing this and test this feature here.
         custom_ast = parse("{ a }")
 
         class CustomFieldNode(SelectionNode):
@@ -1129,6 +1167,7 @@ def allows_skipping_different_sub_trees(skip_action):
 
         class TestVisitor(Visitor):
             def __init__(self, name):
+                super().__init__()
                 self.name = name
 
             def enter(self, *args):
@@ -1232,6 +1271,7 @@ def allows_early_exit_from_different_points(break_action):
 
         class TestVisitor(Visitor):
             def __init__(self, name):
+                super().__init__()
                 self.name = name
 
             def enter(self, *args):
@@ -1323,6 +1363,7 @@ def allows_early_exit_from_leaving_different_points(break_action):
 
         class TestVisitor(Visitor):
             def __init__(self, name):
+                super().__init__()
                 self.name = name
 
             def enter(self, *args):

tests/utilities/test_type_info.py

@@ -101,6 +101,7 @@ def supports_different_operation_types():
 
         class TestVisitor(Visitor):
             def __init__(self):
+                super().__init__()
                 self.root_types = {}
 
             def enter_operation_definition(self, node: OperationDefinitionNode, *_args):
@@ -121,6 +122,7 @@ def provide_exact_same_arguments_to_wrapped_visitor():
 
         class TestVisitor(Visitor):
             def __init__(self):
+                super().__init__()
                 self.args = []
 
             def enter(self, *args):