diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..7644146f --- /dev/null +++ b/.jules/bolt.md @@ -0,0 +1,3 @@ +## 2025-02-23 - Explicit Stack AST Traversal +**Learning:** Recursive `yield from` in hot-path AST traversals introduces significant overhead. Precomputing reversed children via `node._fields` and using an explicit stack preserves traversal order and lazy evaluation while being faster. +**Action:** Use an explicit stack combined with `yield` and `reversed()` children list instead of `yield from` recursion for hot-path AST traversal. diff --git a/src/wardline/scanner/ast_primitives.py b/src/wardline/scanner/ast_primitives.py index 70f565b3..5465efe3 100644 --- a/src/wardline/scanner/ast_primitives.py +++ b/src/wardline/scanner/ast_primitives.py @@ -93,6 +93,25 @@ def build_import_alias_map( return alias_map +def _reversed_children(node: ast.AST) -> list[ast.AST]: + """Return a list of AST children in reverse order, extracting directly from _fields + to avoid ast.iter_child_nodes overhead in hot paths.""" + children: list[ast.AST] = [] + for field in node._fields: + try: + value = getattr(node, field) + except AttributeError: + continue + if isinstance(value, list): + for item in value: + if isinstance(item, ast.AST): + children.append(item) + elif isinstance(value, ast.AST): + children.append(value) + children.reverse() + return children + + def iter_calls_in_function_body( node: ast.FunctionDef | ast.AsyncFunctionDef, ) -> Iterator[ast.Call]: @@ -104,39 +123,39 @@ 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``. """ + stack: list[ast.AST] = list(reversed(node.body)) + + while stack: + current = stack.pop() - 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) - return + args = current.args + for kw_default in reversed(args.kw_defaults): + if kw_default is not None: + stack.append(kw_default) + stack.extend(reversed(args.defaults)) + stack.extend(reversed(current.decorator_list)) + continue + 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) - return + for keyword in reversed(current.keywords): + stack.append(keyword.value) + stack.extend(reversed(current.bases)) + stack.extend(reversed(current.decorator_list)) + continue + if isinstance(current, ast.Lambda): - yield from _walk_argument_defaults(current.args) - return + args = current.args + for kw_default in reversed(args.kw_defaults): + if kw_default is not None: + stack.append(kw_default) + stack.extend(reversed(args.defaults)) + continue + if isinstance(current, ast.Call): yield current - for child in ast.iter_child_nodes(current): - yield from walk_node(child) - - def _walk_argument_defaults(args: ast.arguments) -> Iterator[ast.Call]: - for default in args.defaults: - yield from walk_node(default) - for kw_default in args.kw_defaults: - if kw_default is None: - continue - yield from walk_node(kw_default) - for stmt in node.body: - yield from walk_node(stmt) + stack.extend(_reversed_children(current)) def resolve_self_method_fqn( diff --git a/src/wardline/scanner/rules/_ast_helpers.py b/src/wardline/scanner/rules/_ast_helpers.py index 7c3b52ff..3f95bbe3 100644 --- a/src/wardline/scanner/rules/_ast_helpers.py +++ b/src/wardline/scanner/rules/_ast_helpers.py @@ -33,15 +33,36 @@ _RAISING_CONVERSION_NAMES: frozenset[str] = frozenset({"int", "float", "complex", "Decimal", "Fraction", "UUID"}) +def _reversed_children(node: ast.AST) -> list[ast.AST]: + """Return a list of AST children in reverse order, extracting directly from _fields + to avoid ast.iter_child_nodes overhead in hot paths.""" + children: list[ast.AST] = [] + for field in node._fields: + try: + value = getattr(node, field) + except AttributeError: + continue + if isinstance(value, list): + for item in value: + if isinstance(item, ast.AST): + children.append(item) + elif isinstance(value, ast.AST): + children.append(value) + children.reverse() + return children + + 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)): + stack = _reversed_children(node) + while stack: + current = stack.pop() + if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): continue - if isinstance(child, ast.stmt): - yield child - yield from _own_statements(child) + if isinstance(current, ast.stmt): + yield current + stack.extend(_reversed_children(current)) def _own_reachable_statements( @@ -67,14 +88,16 @@ 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): - if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): - yield child - elif isinstance(child, ast.stmt): + stack = _reversed_children(node) + while stack: + current = stack.pop() + if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): + yield current + elif isinstance(current, ast.stmt): continue else: - yield child - yield from _walk_own_non_stmt_children(child) + yield current + stack.extend(_reversed_children(current)) def _reachable_statements_in_block( @@ -639,9 +662,11 @@ 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): - if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): - yield child + stack = _reversed_children(node) + while stack: + current = stack.pop() + if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): + yield current else: - yield child - yield from _walk_own(child) + yield current + stack.extend(_reversed_children(current))