|
5 | 5 | import random
|
6 | 6 | import warnings
|
7 | 7 | from _ast import AsyncFunctionDef, ClassDef, FunctionDef
|
8 |
| -from collections import defaultdict |
| 8 | +from collections import defaultdict, deque |
9 | 9 | from functools import cache
|
10 | 10 | from pathlib import Path
|
11 | 11 | from typing import TYPE_CHECKING, Any, Optional
|
@@ -602,13 +602,27 @@ def filter_files_optimized(file_path: Path, tests_root: Path, ignore_paths: list
|
602 | 602 |
|
603 | 603 |
|
604 | 604 | def function_has_return_statement(function_node: FunctionDef | AsyncFunctionDef) -> bool:
|
605 |
| - # Custom DFS, return True as soon as a Return node is found |
606 |
| - stack = [function_node] |
| 605 | + # Custom DFS, return True as soon as a Return node is found (optimized) |
| 606 | + stack = deque([function_node]) |
| 607 | + stack_pop = stack.pop |
| 608 | + stack_extend = stack.extend |
| 609 | + |
| 610 | + # Inline version of ast.iter_child_nodes, avoids generator/function call overhead |
| 611 | + def fast_iter_child_nodes(node): |
| 612 | + for name in node._fields: |
| 613 | + value = getattr(node, name, None) |
| 614 | + if isinstance(value, list): |
| 615 | + for item in value: |
| 616 | + if isinstance(item, ast.AST): |
| 617 | + yield item |
| 618 | + elif isinstance(value, ast.AST): |
| 619 | + yield value |
| 620 | + |
607 | 621 | while stack:
|
608 |
| - node = stack.pop() |
| 622 | + node = stack_pop() |
609 | 623 | if isinstance(node, ast.Return):
|
610 | 624 | return True
|
611 |
| - stack.extend(ast.iter_child_nodes(node)) |
| 625 | + stack_extend(fast_iter_child_nodes(node)) |
612 | 626 | return False
|
613 | 627 |
|
614 | 628 |
|
|
0 commit comments