⚡️ Speed up function collect_existing_class_names by 188% in PR #1660 (unstructured-inference)#1838
Conversation
The optimization replaces `ast.walk(tree)` — which visits every node in the AST — with a manual stack-based traversal that only descends into container node types (`Module`, `ClassDef`, `FunctionDef`, control-flow statements, etc.) where `ClassDef` nodes can appear. This eliminates traversal of leaf nodes like `Name`, `Constant`, `Load`, and `Store`, which constitute the bulk of an AST but never contain class definitions. The profiler shows the original single-line comprehension spent 100% of runtime (117.7 ms) in `ast.walk`, while the optimized version completes in 36.1 ms (3.26× faster) by skipping ~60–80% of nodes depending on AST density. Tests confirm correctness across nested classes, control-flow scopes, and large trees with 1000+ classes.
|
Claude finished @misrasaurabh1's task in 4m 51s —— View job PR Review Summary
Prek Checks✅ No linting issues. Mypy reports no type errors. Code Review
The selective traversal lists node types that can contain
The original
Three functions (
Duplicate DetectionNo duplicates detected. The changes are self-contained within a single file. Test CoverageNo unit tests exist for the modified helper functions ( Bot PR StatusPR #1840 ( Last updated: 2026-03-16T20:02Z |
The optimization replaced `any()` generator expressions with explicit early-return for-loops in four helper functions (`_is_namedtuple_class`, `_class_has_explicit_init`, `_has_descriptor_like_class_fields`, and `_has_non_property_method_decorator`), eliminating the overhead of building generator objects and calling the `any()` builtin. Line profiler data shows `_class_has_explicit_init` dropped from 1.85 ms to 0.96 ms (48% faster), and `_is_namedtuple_class` improved from 97 µs to 53 µs (46% faster), because the optimized code avoids allocating iterator state and returns immediately upon finding a match instead of completing the generator. The 51% overall runtime improvement (1.43 ms → 948 µs) comes from these cumulative reductions in per-call overhead across thousands of invocations during AST traversal. Test suite confirms no behavioral changes across all edge cases including dataclasses, decorators, and size-limit boundaries.
⚡️ Codeflash found optimizations for this PR📄 51% (0.51x) speedup for
|
…2026-03-16T19.35.10 ⚡️ Speed up function `_should_use_raw_project_class_context` by 51% in PR #1838 (`codeflash/optimize-pr1660-2026-03-16T19.13.06`)
|
This PR is now faster! 🚀 @misrasaurabh1 accepted my optimizations from: |
⚡️ This pull request contains optimizations for PR #1660
If you approve this dependent PR, these changes will be merged into the original PR branch
unstructured-inference.📄 188% (1.88x) speedup for
collect_existing_class_namesincodeflash/languages/python/context/code_context_extractor.py⏱️ Runtime :
18.6 milliseconds→6.44 milliseconds(best of5runs)⚡️ This change will improve the performance of the following benchmarks:
🔻 This change will degrade the performance of the following benchmarks:
{benchmark_info_degraded}
📝 Explanation and details
The optimization replaces
ast.walk(tree)— which visits every node in the AST — with a manual stack-based traversal that only descends into container node types (Module,ClassDef,FunctionDef, control-flow statements, etc.) whereClassDefnodes can appear. This eliminates traversal of leaf nodes likeName,Constant,Load, andStore, which constitute the bulk of an AST but never contain class definitions. The profiler shows the original single-line comprehension spent 100% of runtime (117.7 ms) inast.walk, while the optimized version completes in 36.1 ms (3.26× faster) by skipping ~60–80% of nodes depending on AST density. Tests confirm correctness across nested classes, control-flow scopes, and large trees with 1000+ classes.✅ Correctness verification report:
🌀 Click to see Generated Regression Tests
⏪ Click to see Replay Tests
benchmarks/codeflash_replay_tests_xz10rm7d/test_tests_benchmarks_test_benchmark_code_extract_code_context__replay_test_0.py::test_codeflash_languages_python_context_code_context_extractor_collect_existing_class_names_test_benchmark_extractTo edit these changes
git checkout codeflash/optimize-pr1660-2026-03-16T19.13.06and push.