Feedback from Syrus

Co-authored-by: Syrus Akbary <me@syrusakbary.com>

Author: alexchamberlain

src/graphql/language/visitor.py

@@ -200,19 +200,12 @@ def __init_subclass__(cls) -> None:
 
     def get_visit_fn(self, kind: str, is_leaving: bool = False) -> Callable:
         """Get the visit function for the given node kind and direction."""
-
         key = (kind, is_leaving)
-        if key in self._visit_fns:
-            return self._visit_fns[key]
-
-        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)
-
-        self._visit_fns[key] = visit_fn
-
-        return visit_fn
+        if key not in self._visit_fns:
+            method = "leave" if is_leaving else "enter"
+            fn = getattr(self, f"{method}_{kind}", None)
+            self._visit_fns[key] = fn or getattr(self, method, None)
+        return self._visit_fns[key]
 
 
 class Stack(NamedTuple):

src/graphql/utilities/type_info.py

@@ -148,12 +148,10 @@ def leave(self, node: Node) -> None:
 
     def _get_method(self, direction: str, kind: str) -> Optional[Callable[[], None]]:
         key = (direction, kind)
-        if key in self._visit_fns:
-            return self._visit_fns[key]
-
-        fn = getattr(self, f"{direction}_{kind}", None)
-        self._visit_fns[key] = fn
-        return fn
+        if key not in self._visit_fns:
+            fn = getattr(self, f"{direction}_{kind}", None)
+            self._visit_fns[key] = fn
+        return self._visit_fns[key]
 
 
     # noinspection PyUnusedLocal