⚡ Bolt: optimize AST traversal hot-paths

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2024-06-16 - Replace yield from with iterative append in hot-path AST traversals
+**Learning:** In heavily used AST traversal functions like `iter_calls_in_function_body` and `_own_statements`, using `yield from` recursively and `isinstance()` checks creates measurable overhead.
+**Action:** For hot-path AST traversals, use an internal walk function that builds a `list` with `.append()`, and replace `isinstance(child, ast.X)` with `type(child) is ast.X` (using `# type: ignore` where required). Convert the final list to an iterator if the API expects one.

src/wardline/core/autofix.py

@@ -50,12 +50,21 @@ def has_comment_in_span(
 
 
 def _own_statements(node: ast.AST) -> Iterator[ast.stmt]:
-    for child in ast.iter_child_nodes(node):
-        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
-            continue
-        if isinstance(child, ast.stmt):
-            yield child
-        yield from _own_statements(child)
+    # Performance optimization: iterative append with explicit type checks avoids
+    # recursive `yield from` overhead and is significantly faster for hot paths.
+    result: list[ast.stmt] = []
+
+    def walk(current: ast.AST) -> None:
+        for child in ast.iter_child_nodes(current):
+            t = type(child)
+            if t is ast.FunctionDef or t is ast.AsyncFunctionDef or t is ast.ClassDef:
+                continue
+            if isinstance(child, ast.stmt):
+                result.append(child)
+            walk(child)
+
+    walk(node)
+    return iter(result)
 
 
 def get_assert_nodes_for_function(func_node: ast.FunctionDef | ast.AsyncFunctionDef) -> list[ast.Assert]:

src/wardline/scanner/ast_primitives.py

@@ -104,39 +104,45 @@ def iter_calls_in_function_body(
     Header expressions that execute in the enclosing scope (decorators, default
     values, base classes, metaclass keywords) are still attributed to ``node``.
     """
-
-    def walk_node(current: ast.AST) -> Iterator[ast.Call]:
-        if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef)):
-            for decorator in current.decorator_list:
-                yield from walk_node(decorator)
-            yield from _walk_argument_defaults(current.args)
+    # Performance optimization: iterative append with explicit type checks avoids
+    # recursive `yield from` overhead and is significantly faster for hot paths.
+    result: list[ast.Call] = []
+
+    def walk_node(current: ast.AST) -> None:
+        t = type(current)
+        if t is ast.FunctionDef or t is ast.AsyncFunctionDef:
+            for decorator in current.decorator_list:  # type: ignore[attr-defined]
+                walk_node(decorator)
+            _walk_argument_defaults(current.args)  # type: ignore[attr-defined]
             return
-        if isinstance(current, ast.ClassDef):
-            for decorator in current.decorator_list:
-                yield from walk_node(decorator)
-            for base in current.bases:
-                yield from walk_node(base)
-            for keyword in current.keywords:
-                yield from walk_node(keyword.value)
+        if t is ast.ClassDef:
+            for decorator in current.decorator_list:  # type: ignore[attr-defined]
+                walk_node(decorator)
+            for base in current.bases:  # type: ignore[attr-defined]
+                walk_node(base)
+            for keyword in current.keywords:  # type: ignore[attr-defined]
+                walk_node(keyword.value)
             return
-        if isinstance(current, ast.Lambda):
-            yield from _walk_argument_defaults(current.args)
+        if t is ast.Lambda:
+            _walk_argument_defaults(current.args)  # type: ignore[attr-defined]
             return
-        if isinstance(current, ast.Call):
-            yield current
+        if t is ast.Call:
+            result.append(current)  # type: ignore[arg-type]
         for child in ast.iter_child_nodes(current):
-            yield from walk_node(child)
+            walk_node(child)
 
-    def _walk_argument_defaults(args: ast.arguments) -> Iterator[ast.Call]:
+    def _walk_argument_defaults(args: ast.arguments) -> None:
         for default in args.defaults:
-            yield from walk_node(default)
+            walk_node(default)
         for kw_default in args.kw_defaults:
             if kw_default is None:
                 continue
-            yield from walk_node(kw_default)
+            walk_node(kw_default)
 
     for stmt in node.body:
-        yield from walk_node(stmt)
+        walk_node(stmt)
+
+    return iter(result)
 
 
 def resolve_self_method_fqn(

src/wardline/scanner/rules/_ast_helpers.py

@@ -21,12 +21,21 @@
 def _own_statements(node: ast.AST) -> Iterator[ast.stmt]:
     """Yield every statement in *node*'s own scope, not descending into nested
     def/class bodies. Includes the bodies of if/for/while/try/with at any depth."""
-    for child in ast.iter_child_nodes(node):
-        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
-            continue
-        if isinstance(child, ast.stmt):
-            yield child
-        yield from _own_statements(child)
+    # Performance optimization: iterative append with explicit type checks avoids
+    # recursive `yield from` overhead and is significantly faster for hot paths.
+    result: list[ast.stmt] = []
+
+    def walk(current: ast.AST) -> None:
+        for child in ast.iter_child_nodes(current):
+            t = type(child)
+            if t is ast.FunctionDef or t is ast.AsyncFunctionDef or t is ast.ClassDef:
+                continue
+            if isinstance(child, ast.stmt):
+                result.append(child)
+            walk(child)
+
+    walk(node)
+    return iter(result)
 
 
 def own_except_handlers(node: ast.FunctionDef | ast.AsyncFunctionDef) -> Iterator[ast.ExceptHandler]: