diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..a9855dac --- /dev/null +++ b/.jules/bolt.md @@ -0,0 +1,3 @@ +## 2025-02-18 - Optimized AST traversal using explicit stack and manual `_fields` iteration +**Learning:** Python's recursive generator syntax (`yield from`) incurs noticeable overhead in deep or wide ASTs, making it slower and more memory-intensive on hot paths than an explicit stack loop. Additionally, `ast.iter_child_nodes()` overhead can be bypassed by manually iterating over a node's `_fields`. To correctly reverse the traversal order (so stack `pop()` processes items depth-first left-to-right) and avoid eager whole-tree list materialization, the children for a *single node* should be extracted into a list and reversed using `stack.extend(reversed(children))`. Also, `isinstance(val, ast.AST)` is necessary when iterating over `_fields` since some fields hold literal strings or booleans, which lack `_fields` themselves. +**Action:** Use an explicit stack (with `yield`, not `yield from`) for hot-path AST traversals. Manually iterate over `_fields`, safely handle `AttributeError`s for missing fields, collect a single node's children into a list, and reverse that list before extending the stack. diff --git a/src/wardline/scanner/analyzer.py b/src/wardline/scanner/analyzer.py index b3660ad3..ebe51eff 100644 --- a/src/wardline/scanner/analyzer.py +++ b/src/wardline/scanner/analyzer.py @@ -453,16 +453,32 @@ def _bind_call_site_arguments_to_parameters( return result def _iter_l2_body_nodes(node: ast.FunctionDef | ast.AsyncFunctionDef) -> Iterator[ast.AST]: - def walk(current: ast.AST) -> Iterator[ast.AST]: - for child in ast.iter_child_nodes(current): - if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): + stack: list[ast.AST] = [] + + def push_children(current: ast.AST) -> None: + children = [] + for field in current._fields: + try: + val = getattr(current, field) + except AttributeError: continue - yield child - yield from walk(child) + if isinstance(val, list): + for item in val: + if isinstance(item, ast.AST): + children.append(item) + elif isinstance(val, ast.AST): + children.append(val) + stack.extend(reversed(children)) for stmt in node.body: yield stmt - yield from walk(stmt) + push_children(stmt) + while stack: + current = stack.pop() + if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): + continue + yield current + push_children(current) def _assignment_targets(node: ast.AST) -> list[ast.expr]: if isinstance(node, ast.Assign): diff --git a/src/wardline/scanner/ast_primitives.py b/src/wardline/scanner/ast_primitives.py index 70f565b3..dc0d0a2a 100644 --- a/src/wardline/scanner/ast_primitives.py +++ b/src/wardline/scanner/ast_primitives.py @@ -104,39 +104,57 @@ 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)) + + def push_children(current: ast.AST) -> None: + children = [] + for field in current._fields: + try: + val = getattr(current, field) + except AttributeError: + continue + if isinstance(val, list): + for item in val: + if isinstance(item, ast.AST): + children.append(item) + elif isinstance(val, ast.AST): + children.append(val) + stack.extend(reversed(children)) + + 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 + for kw in reversed(current.args.kw_defaults): + if kw is not None: + stack.append(kw) + for default in reversed(current.args.defaults): + stack.append(default) + for decorator in reversed(current.decorator_list): + stack.append(decorator) + 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) + for base in reversed(current.bases): + stack.append(base) + for decorator in reversed(current.decorator_list): + stack.append(decorator) + continue + if isinstance(current, ast.Lambda): - yield from _walk_argument_defaults(current.args) - return + for kw in reversed(current.args.kw_defaults): + if kw is not None: + stack.append(kw) + for default in reversed(current.args.defaults): + stack.append(default) + 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) + push_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..b82bb522 100644 --- a/src/wardline/scanner/rules/_ast_helpers.py +++ b/src/wardline/scanner/rules/_ast_helpers.py @@ -36,12 +36,32 @@ 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(current: ast.AST) -> None: + children = [] + for field in current._fields: + try: + val = getattr(current, field) + except AttributeError: + continue + if isinstance(val, list): + for item in val: + if isinstance(item, ast.AST): + children.append(item) + elif isinstance(val, ast.AST): + children.append(val) + stack.extend(reversed(children)) + + 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 +87,34 @@ 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(current: ast.AST) -> None: + children = [] + for field in current._fields: + try: + val = getattr(current, field) + except AttributeError: + continue + if isinstance(val, list): + for item in val: + if isinstance(item, ast.AST): + children.append(item) + elif isinstance(val, ast.AST): + children.append(val) + stack.extend(reversed(children)) + + 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( @@ -635,13 +675,34 @@ def handler_substitutes_on_failure(handler: ast.ExceptHandler, returned_names: f def own_nodes(node: ast.AST) -> Iterator[ast.AST]: """Yield *node* itself and all descendant nodes in its own scope (skipping nested scopes).""" yield node - yield from _walk_own(node) - -def _walk_own(node: ast.AST) -> Iterator[ast.AST]: - for child in ast.iter_child_nodes(node): + stack: list[ast.AST] = [] + + def push_children(current: ast.AST) -> None: + children = [] + for field in current._fields: + try: + val = getattr(current, field) + except AttributeError: + continue + if isinstance(val, list): + for item in val: + if isinstance(item, ast.AST): + children.append(item) + elif isinstance(val, ast.AST): + children.append(val) + stack.extend(reversed(children)) + + 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) + + +def _walk_own(node: ast.AST) -> Iterator[ast.AST]: # pragma: no cover + raise NotImplementedError("Replaced by inline own_nodes logic")