⚡️ Speed up function collect_existing_class_names by 351% in PR #1498 (cf-simplify-context-extraction)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1498

If you approve this dependent PR, these changes will be merged into the original PR branch cf-simplify-context-extraction.

This PR will be automatically closed if the original PR is merged.

---

📄 351% (3.51x) speedup for collect_existing_class_names in codeflash/languages/python/context/code_context_extractor.py

⏱️ Runtime : 2.36 milliseconds → 523 microseconds (best of 17 runs)

📝 Explanation and details

The optimized code achieves a 350% speedup (2.36ms → 523μs) by replacing the generic ast.walk() traversal with a targeted stack-based iteration that only visits nodes where class definitions can appear.

Key Performance Improvement:

The original implementation uses ast.walk(tree), which performs an exhaustive depth-first traversal of every single node in the AST—including expressions, literals, operators, and other leaf nodes that can never contain class definitions. For a typical Python module, this means checking thousands of irrelevant nodes.

The optimized version uses a stack-based approach that only descends into structural nodes (ClassDef, FunctionDef, If, For, While, With, Try blocks) where classes can actually be defined. This dramatically reduces the number of nodes visited and isinstance() checks performed.

Why This Matters:

From the test results, we see consistent 200-700% speedups across all scenarios:

• Empty modules: 579% faster (5.37μs → 791ns) - minimal traversal overhead
• Simple cases: 200-400% faster - fewer nodes to check
• Complex nested structures: 405% faster (37.2μs → 7.37μs) - targeted descent pays off
• Large modules (500 classes): 280% faster (869μs → 228μs) - scales better
• Mixed workloads: 558% faster (799μs → 121μs) - avoids non-class nodes

Impact on Workloads:

Based on the function references showing this is called from build_testgen_context, this optimization benefits test generation workflows that analyze Python code structure. Since class extraction is likely performed repeatedly during code analysis, the 4x speedup directly improves overall test generation throughput. The optimization is particularly effective for large codebases with many classes and complex nesting patterns, as demonstrated by the benchmark results.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 33 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

from __future__ import annotations

# imports
import ast  # used to build AST Module instances for testing

import pytest  # used for our unit tests
from codeflash.languages.python.context.code_context_extractor import \
    collect_existing_class_names

def test_empty_module_returns_empty_set():
    # Parse an empty source module to get an ast.Module with no ClassDef nodes
    tree = ast.parse("")  # empty module
    # The function should return an empty set for an AST without any class definitions
    codeflash_output = collect_existing_class_names(tree) # 5.95μs -> 822ns (624% faster)

import ast

import pytest
from codeflash.languages.python.context.code_context_extractor import \
    collect_existing_class_names

def test_single_class_definition():
    """Test that a single class definition is correctly collected."""
    code = "class MyClass:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 10.2μs -> 3.21μs (218% faster)

def test_multiple_classes_at_module_level():
    """Test that multiple classes at module level are all collected."""
    code = "class ClassA:\n    pass\nclass ClassB:\n    pass\nclass ClassC:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 12.9μs -> 4.26μs (202% faster)

def test_nested_class_definitions():
    """Test that nested classes are collected along with parent classes."""
    code = "class Outer:\n    class Inner:\n        pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 9.96μs -> 2.71μs (267% faster)

def test_class_with_inheritance():
    """Test that class inheritance does not affect class name collection."""
    code = "class Parent:\n    pass\nclass Child(Parent):\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 12.2μs -> 3.46μs (252% faster)

def test_class_with_methods():
    """Test that classes with methods are correctly identified."""
    code = "class MyClass:\n    def method(self):\n        pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 14.0μs -> 2.95μs (374% faster)

def test_class_with_attributes():
    """Test that classes with class attributes are correctly identified."""
    code = "class MyClass:\n    attr = 42"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 11.4μs -> 2.50μs (358% faster)

def test_function_not_collected():
    """Test that function definitions are not collected, only classes."""
    code = "def my_function():\n    pass\nclass MyClass:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 13.0μs -> 3.26μs (299% faster)

def test_deeply_nested_classes():
    """Test that deeply nested classes are all collected."""
    code = "class A:\n    class B:\n        class C:\n            class D:\n                pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 12.4μs -> 3.27μs (280% faster)

def test_empty_module():
    """Test that an empty module returns an empty set."""
    code = ""
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 5.37μs -> 791ns (579% faster)

def test_module_with_only_functions():
    """Test that a module with only functions returns an empty set."""
    code = "def func1():\n    pass\ndef func2():\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 14.7μs -> 3.08μs (377% faster)

def test_module_with_only_statements():
    """Test that a module with only statements returns an empty set."""
    code = "x = 42\ny = 'hello'"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 12.3μs -> 2.29μs (437% faster)

def test_class_with_special_characters_in_name():
    """Test that classes with underscore characters in names are collected."""
    code = "class _PrivateClass:\n    pass\nclass __DunderClass__:\n    pass\nclass Regular_Class:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 12.3μs -> 3.75μs (229% faster)

def test_class_with_numeric_suffix():
    """Test that classes with numeric suffixes are collected."""
    code = "class Class1:\n    pass\nclass Class2:\n    pass\nclass Class123:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 12.0μs -> 3.67μs (227% faster)

def test_class_with_single_letter_name():
    """Test that classes with single letter names are collected."""
    code = "class A:\n    pass\nclass B:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 10.3μs -> 3.15μs (227% faster)

def test_class_in_function_scope():
    """Test that classes defined inside functions are still collected."""
    code = "def outer():\n    class Inner:\n        pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 12.6μs -> 2.62μs (381% faster)

def test_class_in_if_block():
    """Test that classes defined inside if blocks are collected."""
    code = "if True:\n    class MyClass:\n        pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 10.6μs -> 3.20μs (231% faster)

def test_class_in_for_loop():
    """Test that classes defined inside for loops are collected."""
    code = "for i in range(1):\n    class LoopClass:\n        pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 14.6μs -> 3.22μs (353% faster)

def test_class_in_try_except():
    """Test that classes defined inside try-except blocks are collected."""
    code = "try:\n    class TryClass:\n        pass\nexcept:\n    class ExceptClass:\n        pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 13.5μs -> 4.23μs (220% faster)

def test_class_with_decorators():
    """Test that decorated classes are collected."""
    code = "@decorator\nclass DecoratedClass:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 9.78μs -> 2.34μs (317% faster)

def test_class_with_multiple_decorators():
    """Test that classes with multiple decorators are collected."""
    code = "@decorator1\n@decorator2\n@decorator3\nclass DecoratedClass:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 12.0μs -> 2.40μs (401% faster)

def test_duplicate_class_names_returns_single_entry():
    """Test that duplicate class names result in a single set entry."""
    code = "class DuplicateClass:\n    pass\nclass DuplicateClass:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 10.2μs -> 3.31μs (209% faster)

def test_return_type_is_set():
    """Test that the return type is a set."""
    code = "class MyClass:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 8.19μs -> 2.38μs (243% faster)

def test_set_elements_are_strings():
    """Test that all elements in the returned set are strings."""
    code = "class ClassA:\n    pass\nclass ClassB:\n    pass"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 10.2μs -> 3.15μs (224% faster)

def test_class_with_complex_body():
    """Test that classes with complex bodies (multiple methods, attributes) are collected."""
    code = """
class ComplexClass:
    class_attr = 10
    
    def __init__(self):
        self.instance_attr = 20
    
    def method1(self):
        return self.instance_attr
    
    def method2(self):
        x = 5
        return x
    
    @property
    def prop(self):
        return self.instance_attr
"""
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 44.8μs -> 5.40μs (729% faster)

def test_class_with_class_method_and_static_method():
    """Test that classes with classmethods and staticmethods are collected."""
    code = """
class MyClass:
    @classmethod
    def class_method(cls):
        pass
    
    @staticmethod
    def static_method():
        pass
"""
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 20.0μs -> 3.67μs (444% faster)

def test_mixed_functions_and_classes():
    """Test collecting classes from a module with mixed functions and classes."""
    code = """
def func1():
    pass

class Class1:
    pass

def func2():
    pass

class Class2:
    def method(self):
        pass

def func3():
    pass
"""
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 26.3μs -> 5.30μs (397% faster)

def test_sibling_classes_in_nested_scope():
    """Test that sibling classes in a nested scope are both collected."""
    code = """
class Outer:
    class Inner1:
        pass
    class Inner2:
        pass
"""
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 11.6μs -> 3.86μs (201% faster)

def test_many_classes_at_module_level():
    """Test collecting a large number of classes at module level."""
    # Generate code with 100 classes
    class_definitions = "\n".join([f"class Class{i}:\n    pass" for i in range(100)])
    tree = ast.parse(class_definitions)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 181μs -> 49.3μs (268% faster)
    expected = {f"Class{i}" for i in range(100)}

def test_many_nested_classes_within_outer_class():
    """Test collecting many sibling classes within a single outer class."""
    # Generate code with one outer class containing 50 inner classes
    code = "class Outer:\n"
    for i in range(50):
        code += f"    class Inner{i}:\n        pass\n"
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 97.5μs -> 27.8μs (250% faster)
    expected = {"Outer"} | {f"Inner{i}" for i in range(50)}

def test_complex_mixed_structure():
    """Test collecting classes from a complex structure with multiple nesting levels."""
    code = """
class A:
    def method(self):
        pass
    class B:
        class C:
            pass
        def method(self):
            pass

class D:
    pass

def function():
    class E:
        class F:
            pass

class G:
    class H:
        class I:
            class J:
                pass
"""
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 37.2μs -> 7.37μs (405% faster)
    expected = {"A", "B", "C", "D", "E", "F", "G", "H", "I", "J"}

def test_large_number_of_top_level_classes():
    """Test collecting a very large number of top-level classes."""
    # Generate code with 500 top-level classes
    class_definitions = "\n".join([f"class TopLevel{i}:\n    pass" for i in range(500)])
    tree = ast.parse(class_definitions)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 869μs -> 228μs (280% faster)
    expected = {f"TopLevel{i}" for i in range(500)}

def test_performance_with_large_module():
    """Test that the function handles a large module efficiently."""
    # Create a module with a mix of many elements
    code_parts = []
    
    # Add functions
    for i in range(100):
        code_parts.append(f"def func{i}():\n    pass\n")
    
    # Add classes
    for i in range(100):
        code_parts.append(f"class Class{i}:\n    pass\n")
    
    # Add statements
    for i in range(100):
        code_parts.append(f"var{i} = {i}\n")
    
    code = "".join(code_parts)
    tree = ast.parse(code)
    codeflash_output = collect_existing_class_names(tree); result = codeflash_output # 799μs -> 121μs (558% faster)
    
    expected = {f"Class{i}" for i in range(100)}
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1498-2026-02-16T20.53.40 and push.

[claude]

PR Review Summary

Prek Checks

Fixed 3 issues (auto-fixed and pushed in commit 69d32681):

• 2x W293 blank-line-with-whitespace in code_context_extractor.py
• 1x Q000 bad-quotes-inline-string in code_context_extractor.py

Prek now passes cleanly. Mypy reports no issues.

Code Review

1 critical bug found (see inline comment):

The optimized collect_existing_class_names function checks for ast.FunctionDef but not ast.AsyncFunctionDef. In Python's AST, these are sibling classes, not parent-child. This means class definitions inside async def functions will be silently missed by the optimized version, whereas the original ast.walk() approach found them all.

Additional missing node types: ast.AsyncFor, ast.AsyncWith, ast.Match (Python 3.10+).

Test Coverage

This PR changes only codeflash/languages/python/context/code_context_extractor.py (the other files come from the base branch cf-simplify-context-extraction).

File	PR Coverage	Main Coverage	Status
code_context_extractor.py	92%	N/A (new file)	New file with good coverage
unused_definition_remover.py	91%	N/A (new file)	New file with good coverage
config_consts.py	88%	88%	No change
current.py	95%	95%	No change
support.py	51%	54%	Slight decrease
function_optimizer.py	18%	18%	No change
__init__.py	100%	N/A (new file)	Empty init

Notes:

• New files (code_context_extractor.py, unused_definition_remover.py) exceed the 75% coverage threshold
• support.py shows a slight coverage decrease (54% → 51%), likely due to code being moved out to the new context modules
• function_optimizer.py has low coverage (18%) but this is pre-existing and not introduced by this PR
• 8 test failures exist on both PR and main branches (all in test_tracer.py, unrelated to this PR)

---

Last updated: 2026-02-16T21:15:00Z