perf(langserver): increase performance of visitor a little bit more

Author: d-biehl

hatch.toml

@@ -58,6 +58,10 @@ extra-dependencies = [
 python = "3.11"
 post-install-commands = ["pip install -U -e {root:uri}/../robotframework"]
 
+[envs.rfdevel38]
+python = "3.8"
+post-install-commands = ["pip install -U -e {root:uri}/../robotframework"]
+
 [envs.devel]
 python = "3.8"
 

packages/language_server/src/robotcode/language_server/robotframework/utils/async_ast.py

@@ -1,26 +1,38 @@
 import ast
-from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, Type
+from abc import ABC
+from collections import defaultdict
+
+from robot.parsing.model.statements import Statement
+from typing_extensions import Any, AsyncIterator, Callable, Dict, Iterator, Optional, Type, Union
 
 __all__ = ["iter_fields", "iter_child_nodes", "AsyncVisitor"]
 
 
+def _patch_robot() -> None:
+    if hasattr(Statement, "_fields"):
+        Statement._fields = ()
+
+
+_patch_robot()
+
+
 def iter_fields(node: ast.AST) -> Iterator[Any]:
-    """
-    Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields``
-    that is present on *node*.
-    """
     for field in node._fields:
         try:
             yield field, getattr(node, field)
         except AttributeError:
             pass
 
 
+def iter_field_values(node: ast.AST) -> Iterator[Any]:
+    for field in node._fields:
+        try:
+            yield getattr(node, field)
+        except AttributeError:
+            pass
+
+
 def iter_child_nodes(node: ast.AST) -> Iterator[ast.AST]:
-    """
-    Yield all direct child nodes of *node*, that is, all fields that are nodes
-    and all items of fields that are lists of nodes.
-    """
     for _name, field in iter_fields(node):
         if isinstance(field, ast.AST):
             yield field
@@ -46,60 +58,69 @@ async def iter_nodes(node: ast.AST) -> AsyncIterator[ast.AST]:
                 yield n
 
 
-class VisitorFinder:
-    __NOT_SET = object()
+class _NotSet:
+    pass
+
 
-    def __init__(self) -> None:
-        self.__cache: Dict[Type[Any], Optional[Callable[..., Any]]] = {}
+class VisitorFinder(ABC):
+    __NOT_SET = _NotSet()
+    __cls_finder_cache__: Dict[Type[Any], Union[Callable[..., Any], None, _NotSet]]
 
-    def __find_visitor(self, cls: Type[Any]) -> Optional[Callable[..., Any]]:
-        if cls is ast.AST:
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        cls.__cls_finder_cache__ = defaultdict(lambda: cls.__NOT_SET)
+
+    @classmethod
+    def __find_visitor(cls, node_cls: Type[Any]) -> Optional[Callable[..., Any]]:
+        if node_cls is ast.AST:
             return None
-        method_name = "visit_" + cls.__name__
-        if hasattr(self, method_name):
-            method = getattr(self, method_name)
-            if callable(method):
-                return method  # type: ignore
-        for base in cls.__bases__:
-            method = self._find_visitor(base)
+        method_name = "visit_" + node_cls.__name__
+        method = getattr(cls, method_name, None)
+        if callable(method):
+            return method  # type: ignore[no-any-return]
+        for base in node_cls.__bases__:
+            method = cls._find_visitor(base)
             if method:
-                return method  # type: ignore
+                return method
         return None
 
-    def _find_visitor(self, cls: Type[Any]) -> Optional[Callable[..., Any]]:
-        r = self.__cache.get(cls, self.__NOT_SET)
-        if r is self.__NOT_SET:
-            self.__cache[cls] = r = self.__find_visitor(cls)
-        return r  # type: ignore
+    @classmethod
+    def _find_visitor(cls, node_cls: Type[Any]) -> Optional[Callable[..., Any]]:
+        result = cls.__cls_finder_cache__[node_cls]
+        if result is cls.__NOT_SET:
+            result = cls.__cls_finder_cache__[node_cls] = cls.__find_visitor(node_cls)
+        return result  # type: ignore[return-value]
 
 
 class AsyncVisitor(VisitorFinder):
     async def visit(self, node: ast.AST) -> None:
-        visitor = self._find_visitor(type(node)) or self.generic_visit
-        await visitor(node)
+        visitor = self._find_visitor(type(node)) or self.__class__.generic_visit
+        await visitor(self, node)
 
     async def generic_visit(self, node: ast.AST) -> None:
-        """Called if no explicit visitor function exists for a node."""
-        for _, value in iter_fields(node):
-            if isinstance(value, list):
+        for value in iter_field_values(node):
+            if value is None:
+                continue
+            if isinstance(value, ast.AST):
+                await self.visit(value)
+            elif isinstance(value, list):
                 for item in value:
                     if isinstance(item, ast.AST):
                         await self.visit(item)
-            elif isinstance(value, ast.AST):
-                await self.visit(value)
 
 
 class Visitor(VisitorFinder):
     def visit(self, node: ast.AST) -> None:
-        visitor = self._find_visitor(type(node)) or self.generic_visit
-        visitor(node)
+        visitor = self._find_visitor(type(node)) or self.__class__.generic_visit
+        visitor(self, node)
 
     def generic_visit(self, node: ast.AST) -> None:
-        """Called if no explicit visitor function exists for a node."""
-        for field, value in iter_fields(node):
+        for value in iter_field_values(node):
+            if value is None:
+                continue
             if isinstance(value, list):
                 for item in value:
                     if isinstance(item, ast.AST):
                         self.visit(item)
-            elif isinstance(value, ast.AST):
+            else:
                 self.visit(value)

pyproject.toml

@@ -162,7 +162,7 @@ fail_under = 40
 [tool.ruff]
 line-length = 120
 target-version = "py38"
-extend-exclude = ["bundled/libs"]
+extend-exclude = ["bundled/libs", ".hatch"]
 ignore = ["E741", "N805", "N999", "RUF012"]
 select = [
   "E",
@@ -220,15 +220,15 @@ implicit_optional = true
 disallow_untyped_decorators = false
 disallow_subclassing_any = false
 exclude = [
-  '\.mypy_cache',
-  '\.venv',
-  '\.hatch',
-  'build',
-  'dist',
-  'out',
-  'playground',
-  'scripts',
-  'bundled/libs',
+  '^\.mypy_cache/',
+  '^\.venv/',
+  '^\.hatch/',
+  "^build/",
+  "^dist/",
+  "^out/",
+  "^playground/",
+  "^scripts/",
+  "^bundled/libs/",
 ]
 mypy_path = [
   "typings",