Skip to content

feat: extract imported class definitions for testgen context #1014

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
KRRT7 merged 7 commits into from
Jan 8, 2026
Merged

feat: extract imported class definitions for testgen context #1014

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
0312c37
feat: extract imported class definitions for testgen context
KRRT7 Jan 7, 2026
a4f84a8
Merge branch 'main' into feat/extract-imported-class-definitions
KRRT7 Jan 7, 2026
b3f3097
hooks
KRRT7 Jan 7, 2026
4748a75
temp
KRRT7 Jan 7, 2026
c1ac25f
if stack walking didn't find test info, fall back to environment vari…
KRRT7 Jan 7, 2026
40191e7
revert changes here
KRRT7 Jan 7, 2026
fdb1d61
Merge branch 'main' into feat/extract-imported-class-definitions
KRRT7 Jan 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
171 changes: 170 additions & 1 deletion codeflash/context/code_context_extractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,20 +127,48 @@ def get_code_optimization_context(
remove_docstrings=False,
code_context_type=CodeContextType.TESTGEN,
)

# Extract class definitions for imported types from project modules
# This helps the LLM understand class constructors and structure
imported_class_context = get_imported_class_definitions(testgen_context, project_root_path)
if imported_class_context.code_strings:
# Merge imported class definitions into testgen context
testgen_context = CodeStringsMarkdown(
code_strings=testgen_context.code_strings + imported_class_context.code_strings
)

testgen_markdown_code = testgen_context.markdown
testgen_code_token_length = encoded_tokens_len(testgen_markdown_code)
if testgen_code_token_length > testgen_token_limit:
# First try removing docstrings
testgen_context = extract_code_markdown_context_from_files(
helpers_of_fto_dict,
helpers_of_helpers_dict,
project_root_path,
remove_docstrings=True,
code_context_type=CodeContextType.TESTGEN,
)
# Re-extract imported classes (they may still fit)
imported_class_context = get_imported_class_definitions(testgen_context, project_root_path)
if imported_class_context.code_strings:
testgen_context = CodeStringsMarkdown(
code_strings=testgen_context.code_strings + imported_class_context.code_strings
)
testgen_markdown_code = testgen_context.markdown
testgen_code_token_length = encoded_tokens_len(testgen_markdown_code)
if testgen_code_token_length > testgen_token_limit:
raise ValueError("Testgen code context has exceeded token limit, cannot proceed")
# If still over limit, try without imported class definitions
testgen_context = extract_code_markdown_context_from_files(
helpers_of_fto_dict,
helpers_of_helpers_dict,
project_root_path,
remove_docstrings=True,
code_context_type=CodeContextType.TESTGEN,
)
testgen_markdown_code = testgen_context.markdown
testgen_code_token_length = encoded_tokens_len(testgen_markdown_code)
if testgen_code_token_length > testgen_token_limit:
raise ValueError("Testgen code context has exceeded token limit, cannot proceed")
code_hash_context = hashing_code_context.markdown
code_hash = hashlib.sha256(code_hash_context.encode("utf-8")).hexdigest()

Expand Down Expand Up @@ -489,6 +517,147 @@ def get_function_sources_from_jedi(
return file_path_to_function_source, function_source_list


def get_imported_class_definitions(code_context: CodeStringsMarkdown, project_root_path: Path) -> CodeStringsMarkdown:
"""Extract class definitions for imported types from project modules.

This function analyzes the imports in the extracted code context and fetches
class definitions for any classes imported from project modules. This helps
the LLM understand the actual class structure (constructors, methods, inheritance)
rather than just seeing import statements.

Args:
code_context: The already extracted code context containing imports
project_root_path: Root path of the project

Returns:
CodeStringsMarkdown containing class definitions from imported project modules

"""
import jedi

# Collect all code from the context
all_code = "\n".join(cs.code for cs in code_context.code_strings)

# Parse to find import statements
try:
tree = ast.parse(all_code)
except SyntaxError:
return CodeStringsMarkdown(code_strings=[])

# Collect imported names and their source modules
imported_names: dict[str, str] = {} # name -> module_path
for node in ast.walk(tree):
if isinstance(node, ast.ImportFrom) and node.module:
for alias in node.names:
if alias.name != "*":
imported_name = alias.asname if alias.asname else alias.name
imported_names[imported_name] = node.module

if not imported_names:
return CodeStringsMarkdown(code_strings=[])

# Track which classes we've already extracted to avoid duplicates
extracted_classes: set[tuple[Path, str]] = set() # (file_path, class_name)

# Also track what's already defined in the context
existing_definitions: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef):
existing_definitions.add(node.name)

class_code_strings: list[CodeString] = []

for name, module_name in imported_names.items():
# Skip if already defined in context
if name in existing_definitions:
continue

# Try to find the module file using Jedi
try:
# Create a script that imports the module to resolve it
test_code = f"import {module_name}"
script = jedi.Script(test_code, project=jedi.Project(path=project_root_path))
completions = script.goto(1, len(test_code))

if not completions:
continue

module_path = completions[0].module_path
if not module_path:
continue

# Check if this is a project module (not stdlib/third-party)
if not str(module_path).startswith(str(project_root_path) + os.sep):
continue
if path_belongs_to_site_packages(module_path):
continue

# Skip if we've already extracted this class
if (module_path, name) in extracted_classes:
continue

# Parse the module to find the class definition
module_source = module_path.read_text(encoding="utf-8")
module_tree = ast.parse(module_source)

for node in ast.walk(module_tree):
if isinstance(node, ast.ClassDef) and node.name == name:
# Extract the class source code
lines = module_source.split("\n")
class_source = "\n".join(lines[node.lineno - 1 : node.end_lineno])

# Also extract any necessary imports for the class (base classes, type hints)
class_imports = _extract_imports_for_class(module_tree, node, module_source)

full_source = class_imports + "\n\n" + class_source if class_imports else class_source

class_code_strings.append(CodeString(code=full_source, file_path=module_path))
extracted_classes.add((module_path, name))
break

except Exception:
logger.debug(f"Error extracting class definition for {name} from {module_name}")
continue

return CodeStringsMarkdown(code_strings=class_code_strings)


def _extract_imports_for_class(module_tree: ast.Module, class_node: ast.ClassDef, module_source: str) -> str:
"""Extract import statements needed for a class definition.

This extracts imports for base classes and commonly used type annotations.
"""
needed_names: set[str] = set()

# Get base class names
for base in class_node.bases:
if isinstance(base, ast.Name):
needed_names.add(base.id)
elif isinstance(base, ast.Attribute) and isinstance(base.value, ast.Name):
# For things like abc.ABC, we need the module name
needed_names.add(base.value.id)

# Find imports that provide these names
import_lines: list[str] = []
source_lines = module_source.split("\n")

for node in module_tree.body:
if isinstance(node, ast.Import):
for alias in node.names:
name = alias.asname if alias.asname else alias.name.split(".")[0]
if name in needed_names:
import_lines.append(source_lines[node.lineno - 1])
break
elif isinstance(node, ast.ImportFrom):
for alias in node.names:
name = alias.asname if alias.asname else alias.name
if name in needed_names:
import_lines.append(source_lines[node.lineno - 1])
break

Comment on lines +640 to +657
Copy link
Contributor

@codeflash-ai codeflash-ai bot Jan 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 70% (0.70x) speedup for _extract_imports_for_class in codeflash/context/code_context_extractor.py

⏱️ Runtime : 5.37 milliseconds → 3.16 milliseconds (best of 5 runs)

⚡️ This change will improve the performance of the following benchmarks:

Benchmark File :: Function Original Runtime Expected New Runtime Speedup
tests.benchmarks.test_benchmark_code_extract_code_context::test_benchmark_extract 44.3 seconds 44.3 seconds 0.00%

🔻 This change will degrade the performance of the following benchmarks:

{benchmark_info_degraded}

📝 Explanation and details

The optimized code achieves a 70% speedup through three key optimizations that work synergistically:

Key Optimizations

1. Early Exit for Empty Base Classes

if not needed_names:
    return ""

When a class has no base classes (or only built-in bases like object), the function immediately returns without scanning the AST or splitting the source. This optimization shows dramatic gains in test cases like test_no_base_class (91% faster) and test_performance_with_many_unused_imports (3432% faster).

2. Early Termination via Tracking Remaining Names

The optimized version maintains a remaining_names set that shrinks as imports are found:

remaining_names = needed_names.copy()
# ... during iteration ...
if not remaining_names:
    break

This allows the loop to exit as soon as all needed imports are found, rather than scanning the entire AST. For large modules with many imports at the top, this cuts iteration count roughly in half (29,681 → 14,715 hits on the main loop).

3. Deferred Source Splitting

The original code splits module_source into lines immediately, costing ~6% of runtime even when no imports are found. The optimized version:

  • First collects line numbers of matching imports
  • Only splits the source if imports were actually found
  • Uses list comprehension for efficient line extraction
# Only split when needed
if not import_line_numbers:
    return ""
source_lines = module_source.split("\n")
import_lines = [source_lines[line_num] for line_num in import_line_numbers]

Performance Characteristics

Based on the annotated tests, the optimization excels when:

  • Classes have no base classes (91-3432% faster) - early exit avoids all work
  • Many imports but few used (66-182% faster) - early termination stops scanning quickly
  • Large-scale scenarios (42-120% faster) - cumulative effect of all optimizations

The optimization is slightly slower (12-35%) for small, simple cases due to the overhead of copying needed_names and checking remaining_names. However, these cases are extremely fast in absolute terms (microseconds), making the slowdown negligible.

Impact on Workloads

Looking at the function_references, this function is called from get_imported_class_definitions(), which processes every imported class name in a code context. The optimization is particularly beneficial because:

  1. Many classes inherit from built-ins or have no bases → early exit provides major gains
  2. Real-world modules often have many imports at the top → early termination saves significant work
  3. The function is called repeatedly in a loop → per-call savings compound

The optimization transforms this from an O(imports × classes) operation to one that terminates early in most practical cases, making it well-suited for the hot path where it's used to extract class context for code analysis.

✅ Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests ✅ 344 Passed
⏪ Replay Tests ✅ 11 Passed
🔎 Concolic Coverage Tests 🔘 None Found
📊 Tests Coverage 100.0%
🌀 Click to see Generated Regression Tests 
from __future__ import annotations

import ast

# imports
from codeflash.context.code_context_extractor import _extract_imports_for_class

# unit tests

# Basic Test Cases


def test_single_import_for_base_class():
    # Class inherits from Foo, which is imported
    source = "import Foo\nclass Bar(Foo):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 4.23μs -> 5.74μs (26.4% slower)


def test_import_from_for_base_class():
    # Class inherits from Baz, which is imported from a module
    source = "from mod import Baz\nclass Qux(Baz):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.09μs -> 4.62μs (33.0% slower)


def test_multiple_base_classes():
    # Class inherits from multiple bases, both imported
    source = "from mod import Baz\nimport Foo\nclass Qux(Foo, Baz):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 4.05μs -> 5.45μs (25.7% slower)
    lines = result.split("\n")


def test_import_with_asname():
    # Class inherits from Foo, which is imported as Bar
    source = "import Foo as Bar\nclass Baz(Bar):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.96μs -> 4.13μs (28.5% slower)


def test_importfrom_with_asname():
    # Class inherits from Baz, which is imported as Qux
    source = "from mod import Baz as Qux\nclass Foo(Qux):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.95μs -> 4.27μs (30.9% slower)


def test_attribute_base_class():
    # Class inherits from abc.ABC, which is imported
    source = "import abc\nclass MyClass(abc.ABC):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.18μs -> 4.62μs (31.2% slower)


def test_importfrom_attribute_base_class():
    # Class inherits from foo.Bar, foo is imported from somewhere
    source = "from baz import foo\nclass MyClass(foo.Bar):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.17μs -> 4.57μs (30.5% slower)


def test_no_import_needed():
    # Class inherits from built-in object, no import needed
    source = "class MyClass(object):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.14μs -> 2.67μs (19.8% slower)


def test_no_base_class():
    # Class does not inherit from anything
    source = "class MyClass:\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 1.89μs -> 988ns (90.8% faster)


# Edge Test Cases


def test_import_not_included_if_not_base():
    # Foo is imported, but not used as base class
    source = "import Foo\nclass Bar:\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.59μs -> 917ns (182% faster)


def test_multiple_imports_with_similar_names():
    # Both Foo and FooBar imported, only Foo used as base
    source = "import Foo\nimport FooBar\nclass Bar(Foo):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.53μs -> 4.56μs (22.7% slower)


def test_importfrom_multiple_names():
    # from mod import Foo, Bar; only Foo used as base
    source = "from mod import Foo, Bar\nclass Baz(Foo):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.79μs -> 4.30μs (35.1% slower)


def test_importfrom_with_asname_and_normal():
    # from mod import Foo as Bar, Baz; Bar used as base
    source = "from mod import Foo as Bar, Baz\nclass Qux(Bar):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.87μs -> 4.28μs (33.0% slower)


def test_import_with_dotted_name():
    # import foo.bar; base class is bar (should not match)
    source = "import foo.bar\nclass Baz(bar):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.83μs -> 3.23μs (12.3% slower)


def test_importfrom_star():
    # from mod import *; base class is Foo (should not match)
    source = "from mod import *\nclass Baz(Foo):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.72μs -> 2.63μs (3.50% faster)


def test_import_with_comment():
    # Import line has a comment, should still match
    source = "import Foo  # comment here\nclass Bar(Foo):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.15μs -> 4.16μs (24.1% slower)


def test_importfrom_with_comment():
    # ImportFrom line has a comment
    source = "from mod import Foo  # another comment\nclass Bar(Foo):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.86μs -> 4.12μs (30.6% slower)


def test_base_class_is_attribute_of_attribute():
    # Class inherits from foo.bar.Baz, only foo imported
    source = "import foo\nclass MyClass(foo.bar.Baz):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    # Should extract import for 'foo'
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.80μs -> 1.18μs (139% faster)


def test_base_class_is_attribute_of_attribute_from_import():
    # Class inherits from foo.bar.Baz, foo imported from somewhere
    source = "from baz import foo\nclass MyClass(foo.bar.Baz):\n    pass\n"
    tree = ast.parse(source)
    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef))
    # Should extract import for 'foo'
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.71μs -> 1.32μs (106% faster)


# Large Scale Test Cases


def test_many_imports_and_classes():
    # Many imports and classes, only some imports relevant
    imports = [f"import Foo{i}" for i in range(50)]
    classes = [f"class Bar{i}(Foo{i}):\n    pass" for i in range(50)]
    source = "\n".join(imports + classes)
    tree = ast.parse(source)
    for i in range(50):
        class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == f"Bar{i}")
        codeflash_output = _extract_imports_for_class(tree, class_node, source)
        result = codeflash_output  # 915μs -> 529μs (73.1% faster)


def test_many_importfrom_and_classes():
    # Many from-imports and classes, only some imports relevant
    imports = [f"from mod{i} import Baz{i}" for i in range(50)]
    classes = [f"class Qux{i}(Baz{i}):\n    pass" for i in range(50)]
    source = "\n".join(imports + classes)
    tree = ast.parse(source)
    for i in range(50):
        class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == f"Qux{i}")
        codeflash_output = _extract_imports_for_class(tree, class_node, source)
        result = codeflash_output  # 873μs -> 510μs (71.0% faster)


def test_large_mixed_imports_and_classes():
    # Mix of import, importfrom, asname, and many classes
    imports = []
    classes = []
    for i in range(25):
        imports.append(f"import Foo{i} as Bar{i}")
        classes.append(f"class Baz{i}(Bar{i}):\n    pass")
    for i in range(25, 50):
        imports.append(f"from mod{i} import Baz{i} as Qux{i}")
        classes.append(f"class Foo{i}(Qux{i}):\n    pass")
    source = "\n".join(imports + classes)
    tree = ast.parse(source)
    for i in range(25):
        class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == f"Baz{i}")
        codeflash_output = _extract_imports_for_class(tree, class_node, source)
        result = codeflash_output  # 432μs -> 196μs (120% faster)
    for i in range(25, 50):
        class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == f"Foo{i}")
        codeflash_output = _extract_imports_for_class(tree, class_node, source)
        result = codeflash_output  # 421μs -> 296μs (42.0% faster)


def test_large_scale_attribute_base_classes():
    # Many classes inheriting from mod.FooX, mod imported via from-import
    imports = ["from mod import mod"]
    classes = [f"class Bar{i}(mod.Foo{i}):\n    pass" for i in range(50)]
    source = "\n".join(imports + classes)
    tree = ast.parse(source)
    for i in range(50):
        class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == f"Bar{i}")
        codeflash_output = _extract_imports_for_class(tree, class_node, source)
        result = codeflash_output  # 371μs -> 215μs (72.5% faster)


def test_large_scale_no_imports_needed():
    # Many classes inheriting from 'object', no imports needed
    source = "\n".join([f"class Foo{i}(object):\n    pass" for i in range(100)])
    tree = ast.parse(source)
    for i in range(100):
        class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == f"Foo{i}")
        codeflash_output = _extract_imports_for_class(tree, class_node, source)
        result = codeflash_output  # 1.28ms -> 809μs (57.6% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
import ast  # used to create AST nodes for testing

# imports
from codeflash.context.code_context_extractor import _extract_imports_for_class

# Basic Test Cases


def test_single_base_class_with_simple_import():
    """Test extracting import for a class with one base class from simple import."""
    # Create source code with import and class definition
    source = "import BaseClass\n\nclass MyClass(BaseClass):\n    pass"
    # Parse the source into an AST
    tree = ast.parse(source)
    # Get the class node (second item in body after import)
    class_node = tree.body[1]

    # Call function and verify it extracts the import
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.21μs -> 4.14μs (22.4% slower)


def test_single_base_class_with_from_import():
    """Test extracting import for a class with base class from 'from...import' statement."""
    # Create source with from-import statement
    source = "from module import BaseClass\n\nclass MyClass(BaseClass):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Verify the from-import is extracted
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.00μs -> 4.27μs (29.8% slower)


def test_multiple_base_classes_different_imports():
    """Test extracting multiple imports for class with multiple base classes."""
    # Create source with two imports and class inheriting from both
    source = "import Base1\nfrom module import Base2\n\nclass MyClass(Base1, Base2):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[2]

    # Both imports should be extracted
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.94μs -> 5.37μs (26.5% slower)


def test_no_base_classes():
    """Test that no imports are extracted when class has no base classes."""
    # Create source with import but class doesn't use it
    source = "import SomeModule\n\nclass MyClass:\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Should return empty string since no base classes
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.40μs -> 978ns (146% faster)


def test_unused_imports_not_extracted():
    """Test that imports not used by base classes are not extracted."""
    # Create source with multiple imports but only one used
    source = "import Unused1\nimport Unused2\nfrom module import BaseClass\n\nclass MyClass(BaseClass):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[3]

    # Only the used import should be extracted
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.89μs -> 4.86μs (20.0% slower)


# Edge Test Cases


def test_attribute_base_class():
    """Test extracting import for base class specified as attribute (e.g., abc.ABC)."""
    # Create source with module import and class using module.Class syntax
    source = "import abc\n\nclass MyClass(abc.ABC):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Should extract the module import
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.20μs -> 4.36μs (26.6% slower)


def test_import_with_alias():
    """Test extracting import when base class is imported with alias."""
    # Create source with aliased import
    source = "import original_module as om\n\nclass MyClass(om):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Should extract the aliased import
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.94μs -> 3.91μs (24.7% slower)


def test_from_import_with_alias():
    """Test extracting from-import when base class has alias."""
    # Create source with aliased from-import
    source = "from module import OriginalClass as OC\n\nclass MyClass(OC):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Should extract the aliased from-import
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.08μs -> 4.10μs (25.0% slower)


def test_empty_module():
    """Test handling of empty module with no imports."""
    # Create minimal source with just a class
    source = "class MyClass:\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[0]

    # Should return empty string
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 1.89μs -> 1.03μs (82.8% faster)


def test_base_class_without_import():
    """Test class with base class that has no corresponding import."""
    # Create source where base class is not imported (might be builtin or error)
    source = "class MyClass(SomeClass):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[0]

    # Should return empty string since no matching import
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.10μs -> 2.28μs (7.94% slower)


def test_dotted_import_name():
    """Test import with dotted name (e.g., import os.path)."""
    # Create source with dotted import
    source = "import os.path\n\nclass MyClass(os):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Should extract the import, matching on first part of dotted name
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.21μs -> 4.59μs (29.9% slower)


def test_multiple_names_in_single_import():
    """Test import statement with multiple names where only one is used."""
    # Create source with comma-separated imports
    source = "from module import Class1, Class2, Class3\n\nclass MyClass(Class2):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Should extract the entire import line
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.27μs -> 4.38μs (25.3% slower)


def test_nested_attribute_base_class():
    """Test base class with nested attributes (e.g., module.submodule.Class)."""
    # Create source with import and nested attribute usage
    source = "import collections.abc\n\nclass MyClass(collections.abc):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Should extract import matching the first part of attribute chain
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.31μs -> 4.65μs (28.7% slower)


def test_multiple_classes_same_import():
    """Test that same import is extracted for different classes using it."""
    # Create source with one import and two classes using it
    source = "from module import Base\n\nclass Class1(Base):\n    pass\n\nclass Class2(Base):\n    pass"
    tree = ast.parse(source)

    # Test first class
    class_node1 = tree.body[1]
    codeflash_output = _extract_imports_for_class(tree, class_node1, source)
    result1 = codeflash_output  # 3.17μs -> 4.37μs (27.4% slower)

    # Test second class
    class_node2 = tree.body[2]
    codeflash_output = _extract_imports_for_class(tree, class_node2, source)
    result2 = codeflash_output  # 1.44μs -> 1.92μs (25.0% slower)


def test_multiline_source_correct_line_extraction():
    """Test that correct line is extracted from multiline source."""
    # Create source with multiple lines and imports at different positions
    source = "# Comment\nimport Base1\n\nfrom module import Base2\n\nclass MyClass(Base2):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[2]  # Third statement (after two imports)

    # Should extract the correct line
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 3.48μs -> 4.70μs (26.0% slower)


def test_import_star():
    """Test handling of 'from module import *' statement."""
    # Create source with star import
    source = "from module import *\n\nclass MyClass(SomeBase):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Star imports have special handling - check if extracted
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.77μs -> 2.63μs (5.16% faster)


def test_relative_import():
    """Test handling of relative imports (from . import or from .. import)."""
    # Create source with relative import
    source = "from .module import BaseClass\n\nclass MyClass(BaseClass):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Should extract relative import
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.80μs -> 4.11μs (31.9% slower)


def test_base_class_as_call_not_name():
    """Test that base classes specified as calls (not names) don't cause errors."""
    # Create source where base class is a function call (e.g., Generic[T])
    source = "from typing import Generic\n\nclass MyClass(Generic[int]):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[1]

    # Function calls in bases should be handled gracefully (not extracted)
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 2.56μs -> 1.23μs (109% faster)


# Large Scale Test Cases


def test_many_base_classes():
    """Test class with many base classes requiring many imports."""
    # Create source with 50 different base classes
    num_bases = 50
    imports = [f"from module{i} import Base{i}" for i in range(num_bases)]
    bases = ", ".join([f"Base{i}" for i in range(num_bases)])
    source = "\n".join(imports) + f"\n\nclass MyClass({bases}):\n    pass"

    tree = ast.parse(source)
    class_node = tree.body[num_bases]  # After all imports

    # All imports should be extracted
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 25.1μs -> 29.6μs (15.2% slower)
    result_lines = result.split("\n")
    for i in range(num_bases):
        pass


def test_many_imports_in_module():
    """Test module with many imports where only few are used."""
    # Create source with 100 imports but only 3 used
    num_imports = 100
    imports = [f"from module{i} import Class{i}" for i in range(num_imports)]
    source = "\n".join(imports) + "\n\nclass MyClass(Class10, Class50, Class99):\n    pass"

    tree = ast.parse(source)
    class_node = tree.body[num_imports]

    # Only 3 imports should be extracted
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 29.3μs -> 32.2μs (9.27% slower)
    result_lines = result.split("\n")


def test_large_source_file():
    """Test extraction from large source file with many lines."""
    # Create source with 500 lines including comments, blank lines, and code
    lines = []
    lines.append("import BaseClass")
    for i in range(498):
        lines.append(f"# Comment line {i}")
    lines.append("class MyClass(BaseClass):\n    pass")
    source = "\n".join(lines)

    tree = ast.parse(source)
    # Find the class node (should be last statement)
    class_node = [node for node in tree.body if isinstance(node, ast.ClassDef)][0]

    # Should correctly extract import from line 0
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 38.8μs -> 23.3μs (66.9% faster)


def test_complex_inheritance_hierarchy():
    """Test class with complex mix of import types and base classes."""
    # Create source with various import styles
    source = """import abc
import typing
from collections.abc import Mapping
from typing import Generic, TypeVar
from custom import CustomBase as CB

T = TypeVar('T')

class MyClass(abc.ABC, Mapping, CB):
    pass"""

    tree = ast.parse(source)
    # Find the class node
    class_node = [node for node in tree.body if isinstance(node, ast.ClassDef)][0]

    # Should extract all relevant imports
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 5.91μs -> 7.53μs (21.5% slower)


def test_performance_with_many_unused_imports():
    """Test performance when module has many imports but class uses none."""
    # Create source with 200 imports but class has no bases
    num_imports = 200
    imports = [f"import module{i}" for i in range(num_imports)]
    source = "\n".join(imports) + "\n\nclass MyClass:\n    pass"

    tree = ast.parse(source)
    class_node = tree.body[num_imports]

    # Should quickly return empty string
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 50.7μs -> 1.44μs (3432% faster)


def test_mixed_import_styles_large_scale():
    """Test large module with mixed import styles (import, from-import, aliases)."""
    # Create source with 100 mixed-style imports
    imports = []
    for i in range(100):
        if i % 3 == 0:
            imports.append(f"import module{i}")
        elif i % 3 == 1:
            imports.append(f"from package{i} import Class{i}")
        else:
            imports.append(f"from package{i} import Class{i} as C{i}")

    # Use every 10th import as base class
    bases = []
    expected_imports = []
    for i in range(0, 100, 10):
        if i % 3 == 0:
            bases.append(f"module{i}")
            expected_imports.append(f"import module{i}")
        elif i % 3 == 1:
            bases.append(f"Class{i}")
            expected_imports.append(f"from package{i} import Class{i}")
        else:
            bases.append(f"C{i}")
            expected_imports.append(f"from package{i} import Class{i} as C{i}")

    source = "\n".join(imports) + f"\n\nclass MyClass({', '.join(bases)}):\n    pass"
    tree = ast.parse(source)
    class_node = tree.body[100]

    # Should extract all 10 used imports
    codeflash_output = _extract_imports_for_class(tree, class_node, source)
    result = codeflash_output  # 32.0μs -> 36.6μs (12.7% slower)
    result_lines = result.split("\n")
    for expected in expected_imports:
        pass


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
⏪ Click to see Replay Tests
Test File::Test Function Original ⏱️ Optimized ⏱️ Speedup
benchmarks/codeflash_replay_tests_fuodcj9h/test_tests_benchmarks_test_benchmark_code_extract_code_context__replay_test_0.py::test_codeflash_context_code_context_extractor__extract_imports_for_class_test_benchmark_extract 786μs 332μs 136%✅

To test or edit this optimization locally git merge codeflash/optimize-pr1014-2026-01-07T21.53.54

Click to see suggested changes
Suggested change
# Find imports that provide these names
import_lines: list[str] = []
source_lines = module_source.split("\n")
for node in module_tree.body:
if isinstance(node, ast.Import):
for alias in node.names:
name = alias.asname if alias.asname else alias.name.split(".")[0]
if name in needed_names:
import_lines.append(source_lines[node.lineno - 1])
break
elif isinstance(node, ast.ImportFrom):
for alias in node.names:
name = alias.asname if alias.asname else alias.name
if name in needed_names:
import_lines.append(source_lines[node.lineno - 1])
break
if not needed_names:
return ""
# Find imports that provide these names
import_line_numbers: list[int] = []
remaining_names = needed_names.copy()
for node in module_tree.body:
if not remaining_names:
break
if isinstance(node, ast.Import):
for alias in node.names:
name = alias.asname if alias.asname else alias.name.split(".")[0]
if name in remaining_names:
import_line_numbers.append(node.lineno - 1)
remaining_names.discard(name)
break
elif isinstance(node, ast.ImportFrom):
for alias in node.names:
name = alias.asname if alias.asname else alias.name
if name in remaining_names:
import_line_numbers.append(node.lineno - 1)
remaining_names.discard(name)
break
if not import_line_numbers:
return ""
source_lines = module_source.split("\n")
import_lines: list[str] = [source_lines[line_num] for line_num in import_line_numbers]

return "\n".join(import_lines)


def is_dunder_method(name: str) -> bool:
return len(name) > 4 and name.isascii() and name.startswith("__") and name.endswith("__")

Expand Down
8 changes: 5 additions & 3 deletions codeflash/models/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -655,16 +655,18 @@ def find_func_in_class(self, class_node: cst.ClassDef, func_name: str) -> Option
def get_src_code(self, test_path: Path) -> Optional[str]:
if not test_path.exists():
return None
test_src = test_path.read_text(encoding="utf-8")
module_node = cst.parse_module(test_src)
try:
test_src = test_path.read_text(encoding="utf-8")
module_node = cst.parse_module(test_src)
except Exception:
return None

if self.test_class_name:
for stmt in module_node.body:
if isinstance(stmt, cst.ClassDef) and stmt.name.value == self.test_class_name:
func_node = self.find_func_in_class(stmt, self.test_function_name)
if func_node:
return module_node.code_for_node(func_node).strip()
# class not found
return None

# Otherwise, look for a top level function
Expand Down
11 changes: 11 additions & 0 deletions codeflash/verification/codeflash_capture.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,17 @@ def get_test_info_from_stack(tests_root: str) -> tuple[str, str | None, str, str
# Go to the previous frame
frame = frame.f_back

# If stack walking didn't find test info, fall back to environment variables
if not test_name:
env_test_function = os.environ.get("CODEFLASH_TEST_FUNCTION")
if env_test_function:
test_name = env_test_function
if not test_module_name:
test_module_name = os.environ.get("CODEFLASH_TEST_MODULE", "")
if not test_class_name:
env_class = os.environ.get("CODEFLASH_TEST_CLASS")
test_class_name = env_class if env_class else None

return test_module_name, test_class_name, test_name, line_id


Expand Down
Loading
Loading