From 2521d069a238aa2805b0fdf6fe0997c3e699697e Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 30 Jun 2026 16:39:37 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Optimize=20AST=20traversal?= =?UTF-8?q?=20with=20an=20explicit=20stack?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces recursive `yield from` AST traversal with a flat stack-based approach in several hot-path functions, improving parsing speed. Co-authored-by: tachyon-beep <544926+tachyon-beep@users.noreply.github.com> --- .jules/bolt.md | 3 ++ src/wardline/core/autofix.py | 19 +++++++- src/wardline/scanner/rules/_ast_helpers.py | 57 +++++++++++++++++++--- 3 files changed, 71 insertions(+), 8 deletions(-) create mode 100644 .jules/bolt.md diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..4cd2724b --- /dev/null +++ b/.jules/bolt.md @@ -0,0 +1,3 @@ +## 2024-06-30 - AST Traversal performance optimization +**Learning:** Python's `ast.iter_child_nodes` coupled with deep generator delegation (`yield from`) can add a major overhead for large or deep abstract syntax trees in performance-critical codebases. While an inline C implementation like `list(ast.iter_child_nodes())` is fast, managing a stack implicitly through nested generator frames creates significant delegation lag. +**Action:** When implementing deep AST walks for short-circuited scans (where lazy eval matters), replace recursive `yield from` with a flat list stack. To avoid eager whole-tree computations, iterate over `reversed(getattr(n, "_fields", ()))` directly to fetch immediate children and push them into the stack, avoiding expensive allocations. diff --git a/src/wardline/core/autofix.py b/src/wardline/core/autofix.py index f8660bd3..b9e8020e 100644 --- a/src/wardline/core/autofix.py +++ b/src/wardline/core/autofix.py @@ -50,12 +50,27 @@ def has_comment_in_span( def _own_statements(node: ast.AST) -> Iterator[ast.stmt]: - for child in ast.iter_child_nodes(node): + stack: list[ast.AST] = [] + + def push_children(n: ast.AST) -> None: + for field in reversed(getattr(n, "_fields", ())): + val = getattr(n, field, None) + if isinstance(val, list): + for item in reversed(val): + if isinstance(item, ast.AST): + stack.append(item) + elif isinstance(val, ast.AST): + stack.append(val) + + push_children(node) + + while stack: + child = stack.pop() if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): continue if isinstance(child, ast.stmt): yield child - yield from _own_statements(child) + push_children(child) def get_assert_nodes_for_function(func_node: ast.FunctionDef | ast.AsyncFunctionDef) -> list[ast.Assert]: diff --git a/src/wardline/scanner/rules/_ast_helpers.py b/src/wardline/scanner/rules/_ast_helpers.py index 7c3b52ff..a9ca9813 100644 --- a/src/wardline/scanner/rules/_ast_helpers.py +++ b/src/wardline/scanner/rules/_ast_helpers.py @@ -36,12 +36,27 @@ 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): + stack: list[ast.AST] = [] + + def push_children(n: ast.AST) -> None: + for field in reversed(getattr(n, "_fields", ())): + val = getattr(n, field, None) + if isinstance(val, list): + for item in reversed(val): + if isinstance(item, ast.AST): + stack.append(item) + elif isinstance(val, ast.AST): + stack.append(val) + + push_children(node) + + while stack: + child = stack.pop() if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): continue if isinstance(child, ast.stmt): yield child - yield from _own_statements(child) + push_children(child) def _own_reachable_statements( @@ -67,14 +82,29 @@ def _own_nodes_in_reachable_stmt(stmt: ast.stmt) -> Iterator[ast.AST]: def _walk_own_non_stmt_children(node: ast.AST) -> Iterator[ast.AST]: - for child in ast.iter_child_nodes(node): + stack: list[ast.AST] = [] + + def push_children(n: ast.AST) -> None: + for field in reversed(getattr(n, "_fields", ())): + val = getattr(n, field, None) + if isinstance(val, list): + for item in reversed(val): + if isinstance(item, ast.AST): + stack.append(item) + elif isinstance(val, ast.AST): + stack.append(val) + + push_children(node) + + while stack: + child = stack.pop() if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): yield child elif isinstance(child, ast.stmt): continue else: yield child - yield from _walk_own_non_stmt_children(child) + push_children(child) def _reachable_statements_in_block( @@ -639,9 +669,24 @@ def own_nodes(node: ast.AST) -> Iterator[ast.AST]: def _walk_own(node: ast.AST) -> Iterator[ast.AST]: - for child in ast.iter_child_nodes(node): + stack: list[ast.AST] = [] + + def push_children(n: ast.AST) -> None: + for field in reversed(getattr(n, "_fields", ())): + val = getattr(n, field, None) + if isinstance(val, list): + for item in reversed(val): + if isinstance(item, ast.AST): + stack.append(item) + elif isinstance(val, ast.AST): + stack.append(val) + + push_children(node) + + while stack: + child = stack.pop() if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): yield child else: yield child - yield from _walk_own(child) + push_children(child)