⚡️ Speed up function extract_imports_for_class by 429% in PR #1339 (coverage-no-files)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1339

If you approve this dependent PR, these changes will be merged into the original PR branch coverage-no-files.

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

---

📄 429% (4.29x) speedup for extract_imports_for_class in codeflash/context/code_context_extractor.py

⏱️ Runtime : 2.33 milliseconds → 441 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 428% runtime speedup (2.33ms → 441μs) by replacing the expensive ast.walk(class_node) traversal with direct iteration over class_node.body.

Key Optimization

Original approach: Used ast.walk(class_node) which recursively visits every node in the AST subtree, including all nested function definitions, their arguments, return types, and deeply nested expression nodes. For a typical class with methods, this traverses ~2500 nodes.

Optimized approach: Iterates only class_node.body, which contains just the direct children of the class (typically 200-400 nodes for the same class). This is sufficient because:

• Type annotations for fields are in class_node.body as ast.AnnAssign nodes
• Field assignments with field() calls are in class_node.body as ast.Assign nodes
• Base classes and decorators are already extracted separately before the loop

The line profiler confirms this: the original's ast.walk() loop consumed 66% of total runtime (12.76ms out of 19.3ms), while the optimized version's direct iteration takes only 2.3% (112μs out of 4.96ms).

Additional Refinement

The optimized code also improves the field() detection by changing from checking ast.Call nodes anywhere in the tree to specifically checking ast.Assign nodes where the value is a Call with a Name func. This more accurately targets dataclass field assignments and uses elif to avoid redundant checks.

Test Case Performance

The optimization excels across all test categories:

• Simple classes (2-3 fields): 186-436% faster
• Complex annotations (nested generics): 335-591% faster
• Large-scale tests (50+ fields, 200 imports): 495-949% faster

The performance gain scales with class complexity because larger classes have more nested nodes that ast.walk() unnecessarily traverses, while the optimized version still only iterates the direct body elements.

Impact on Workloads

Based on function_references, extract_imports_for_class is called from:

1. Test suite replay tests - indicating it's in a performance-critical testing path
2. get_code_optimization_context - suggesting it's used during code analysis/optimization workflows

Since the function extracts context for optimization decisions, the 428% speedup directly reduces latency in code analysis pipelines, making the optimization particularly valuable for CI/CD systems or developer tooling that analyzes many classes.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from __future__ import annotations

# imports
import ast
import ast as _ast  # avoid shadowing the module-level ast usage in tests

import pytest  # used for our unit tests
from codeflash.context.code_context_extractor import extract_imports_for_class

# unit tests

def _get_first_class_node(module_source: str) -> tuple[_ast.Module, _ast.ClassDef]:
    """
    Helper to parse source and return (module_tree, first_class_node).
    Raises a helpful assertion if no class found.
    """
    module_tree = ast.parse(module_source)
    for node in module_tree.body:
        if isinstance(node, ast.ClassDef):
            return module_tree, node
    raise AssertionError("No class definition found in test source")

def test_basic_base_and_decorator_imports():
    # Basic scenario: class uses an attribute base (abc.ABC) and a simple decorator (dataclass).
    source = "\n".join([
        "import abc",                     # line 1
        "from dataclasses import dataclass, field",  # line 2
        "",
        "@dataclass",                      # decorator on class
        "class MyBase(abc.ABC):",          # attribute base - should capture 'abc'
        "    x: int = 0",
    ])
    module_tree, class_node = _get_first_class_node(source)
    codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 19.4μs -> 5.01μs (286% faster)

def test_decorator_call_and_attribute_decorator():
    # Ensure decorator calls and attribute-based decorators are handled.
    source = "\n".join([
        "import custommod",                 # should be captured via decorator attribute value
        "from decorators import deco",      # should be captured by decorator name
        "",
        "@custommod.decorator()",           # decorator is a Call with func Attribute(custommod.decorator)
        "@deco()",                          # decorator is a Call with func Name(deco)
        "class C:",                         # no bases
        "    pass",
    ])
    module_tree, class_node = _get_first_class_node(source)
    codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 17.2μs -> 5.23μs (230% faster)

def test_annotations_subscript_tuple_union_and_field_call():
    # Complex annotations: Subscript (List[...] via alias), Tuple, Union (|), and a call to field()
    source = "\n".join([
        "from typing import List as L, Tuple",  # alias L used for subscript, Tuple for tuple annotations
        "from dataclasses import field",        # field() call should be captured
        "from pkg import A, B",                 # A and B referenced in union/tuple
        "",
        "class DataClass:",
        "    a: L[A]                       # List alias used",   # AnnAssign with Subscript
        "    b: Tuple[A, B]                # Tuple annotation",   # AnnAssign with Tuple
        "    c: A | B                      # Union using |",       # AnnAssign: BinOp on annotations
        "    d = field(default=1)          # Call to field (ast.Call with Name func)",
    ])
    module_tree, class_node = _get_first_class_node(source)
    codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 39.8μs -> 8.08μs (393% faster)

    # Expect the three import lines; order is typing (line1), dataclasses (line2), pkg (line3)
    expected = "from typing import List as L, Tuple\nfrom dataclasses import field\nfrom pkg import A, B"

def test_duplicate_imports_and_alias_handling():
    # The same import line can provide multiple names - the importer should be returned once.
    source = "\n".join([
        "from multi import X, Y, Z",  # single import line providing X, Y, Z
        "",
        "class UsesXY:",
        "    a: X",
        "    b: Y",
    ])
    module_tree, class_node = _get_first_class_node(source)
    codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 16.2μs -> 3.75μs (333% faster)

def test_no_matching_imports_returns_empty_string():
    # If the class references names that were not imported, the function returns an empty string.
    source = "\n".join([
        "import unrelated",
        "",
        "class Alone(SomeMissingBase):",
        "    x: UnknownType",
    ])
    module_tree, class_node = _get_first_class_node(source)
    codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 13.9μs -> 3.37μs (312% faster)

def test_attribute_calls_are_not_collected_as_field_names():
    # Calls where the function is an attribute (like module.field()) should NOT be counted as a Name.
    # Only Name-based calls (field()) are collected by the implementation.
    source = "\n".join([
        "import module",                   # present but should NOT be included because only module.field() is used
        "from dataclasses import field",   # field() as bare name would be captured if present
        "",
        "class AttrCall:",
        "    x = module.field()",          # ast.Call with func ast.Attribute -> should NOT add 'module' to needed_names
        "    y = field()",                 # ast.Call with func ast.Name -> should add 'field' to needed_names
    ])
    module_tree, class_node = _get_first_class_node(source)
    codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 19.8μs -> 4.27μs (363% faster)

def test_large_scale_many_imports_and_annotations():
    # Large-scale test: construct many imports and many annotated class attributes
    # Keep the size under 1000 as requested; use 200 as a representative large case.
    n = 200
    lines = []
    # Create import lines: import mod0, import mod1, ...
    for i in range(n):
        lines.append(f"import mod{i}")
    lines.append("")  # blank line before class
    lines.append("class Big:")
    # Create many annotated assignments within the class referencing each mod{i}
    for i in range(n):
        # Each annotation is a simple Name referring to mod{i}
        lines.append(f"    a{i}: mod{i}")
    source = "\n".join(lines)

    module_tree = ast.parse(source)
    # Find the class node explicitly
    class_node = None
    for node in module_tree.body:
        if isinstance(node, ast.ClassDef) and node.name == "Big":
            class_node = node
            break

    codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 737μs -> 123μs (495% faster)

    # The result should contain exactly n import lines, one for each mod{i}, in the same order as in the source.
    result_lines = result.split("\n")

def test_importfrom_with_asname_used_in_annotations():
    # Test that ImportFrom aliases (asname) are honored when matching annotations.
    source = "\n".join([
        "from lib import Thing as T, Other",  # alias T should match annotation using T
        "",
        "class UsesAlias:",
        "    x: T",
        "    y: Other",
    ])
    module_tree, class_node = _get_first_class_node(source)
    codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 16.9μs -> 4.00μs (322% 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

import pytest
from codeflash.context.code_context_extractor import extract_imports_for_class

class TestExtractImportsBasicFunctionality:
    """Test basic functionality of extract_imports_for_class"""

    def test_simple_class_with_no_bases_no_decorators(self):
        """Test extraction from a simple class with no base classes or decorators"""
        source = """
class SimpleClass:
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[0]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 8.74μs -> 2.68μs (226% faster)

    def test_class_with_simple_base_class(self):
        """Test extraction when class inherits from a simple base class"""
        source = """from abc import ABC

class MyClass(ABC):
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 11.7μs -> 3.92μs (198% faster)

    def test_class_with_module_attribute_base(self):
        """Test extraction when class inherits from module.ClassName (attribute access)"""
        source = """import abc

class MyClass(abc.ABC):
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 13.1μs -> 4.09μs (221% faster)

    def test_class_with_single_decorator(self):
        """Test extraction when class has a single decorator"""
        source = """from dataclasses import dataclass

@dataclass
class MyClass:
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 11.3μs -> 3.96μs (186% faster)

    def test_class_with_decorator_call(self):
        """Test extraction when decorator has function call syntax"""
        source = """from functools import lru_cache

@lru_cache(maxsize=128)
class MyClass:
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 14.7μs -> 4.07μs (261% faster)

    def test_class_with_type_annotation_field(self):
        """Test extraction when class has annotated fields"""
        source = """from typing import List

class MyClass:
    items: List[str]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 16.2μs -> 4.20μs (286% faster)

    def test_class_with_multiple_bases(self):
        """Test extraction when class inherits from multiple base classes"""
        source = """from abc import ABC
from collections.abc import Iterable

class MyClass(ABC, Iterable):
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[2]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 13.2μs -> 4.81μs (174% faster)

    def test_class_with_multiple_decorators(self):
        """Test extraction when class has multiple decorators"""
        source = """from dataclasses import dataclass
from functools import total_ordering

@dataclass
@total_ordering
class MyClass:
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[2]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 13.0μs -> 4.74μs (174% faster)

    def test_class_with_aliased_import(self):
        """Test extraction when imports use 'as' aliases"""
        source = """from typing import List as L

class MyClass:
    items: L[str]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 15.8μs -> 4.13μs (283% faster)

class TestExtractImportsEdgeCases:
    """Test edge cases and unusual scenarios"""

    def test_class_with_unused_imports(self):
        """Test that unused imports are not extracted"""
        source = """from typing import List
from os import path

class MyClass:
    items: List[str]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[2]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 16.5μs -> 4.51μs (266% faster)

    def test_empty_class_with_no_annotations(self):
        """Test class with no bases, decorators, or type annotations"""
        source = """class EmptyClass:
    def method(self):
        pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[0]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 12.7μs -> 2.37μs (436% faster)

    def test_class_with_complex_type_annotations(self):
        """Test class with nested generic type annotations"""
        source = """from typing import Dict, List, Optional

class MyClass:
    data: Dict[str, List[Optional[int]]]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 24.0μs -> 5.52μs (335% faster)

    def test_class_with_module_attribute_decorator(self):
        """Test decorator accessed as module.decorator_name"""
        source = """import dataclasses

@dataclasses.dataclass
class MyClass:
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 12.0μs -> 3.41μs (252% faster)

    def test_class_with_union_type_annotation(self):
        """Test class with Union type using pipe syntax (Python 3.10+)"""
        source = """from typing import Union

class MyClass:
    value: Union[int, str]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 18.6μs -> 4.59μs (306% faster)

    def test_class_with_builtin_base_no_import(self):
        """Test class inheriting from builtin with no import needed"""
        source = """class MyClass(Exception):
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[0]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 9.59μs -> 2.83μs (239% faster)

    def test_class_with_duplicate_decorator_usage(self):
        """Test that duplicate imports are not added multiple times"""
        source = """from functools import wraps

@wraps
class MyClass:
    def method(self):
        wraps(lambda: None)
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 21.6μs -> 3.75μs (475% faster)

    def test_class_with_import_star(self):
        """Test handling of 'from module import *' statements"""
        source = """from typing import *

class MyClass:
    items: List
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 12.1μs -> 3.21μs (276% faster)

    def test_class_with_no_bases_attribute_access(self):
        """Test class without base classes doesn't incorrectly match attribute names"""
        source = """import os

class MyClass:
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 8.55μs -> 3.19μs (168% faster)

    def test_class_with_optional_annotation(self):
        """Test class with Optional type annotation"""
        source = """from typing import Optional

class MyClass:
    value: Optional[str]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 16.1μs -> 4.12μs (291% faster)

    def test_class_with_nested_annotation(self):
        """Test class with deeply nested type annotations"""
        source = """from typing import Dict, Tuple, List

class MyClass:
    nested: Dict[str, Tuple[List[int], str]]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 26.9μs -> 5.78μs (366% faster)

    def test_class_with_single_letter_name_not_matching(self):
        """Test that single-letter type variables don't incorrectly match"""
        source = """from typing import TypeVar

T = TypeVar('T')

class MyClass(T):
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[2]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 10.7μs -> 3.57μs (201% faster)

class TestExtractImportsComplexScenarios:
    """Test complex combinations of features"""

    def test_class_with_all_features_combined(self):
        """Test class with bases, decorators, and type annotations all together"""
        source = """from abc import ABC
from dataclasses import dataclass
from typing import List

@dataclass
class MyClass(ABC):
    items: List[str]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[3]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 20.1μs -> 5.89μs (241% faster)

    def test_multiple_classes_in_module(self):
        """Test extracting imports for specific class when multiple exist"""
        source = """from abc import ABC
from typing import List

class ClassA(ABC):
    pass

class ClassB:
    items: List[str]
"""
        module_tree = ast.parse(source)
        class_b = module_tree.body[2]
        codeflash_output = extract_imports_for_class(module_tree, class_b, source); result = codeflash_output # 11.6μs -> 4.61μs (152% faster)

    def test_decorator_with_multiple_arguments(self):
        """Test decorator with complex argument list"""
        source = """from functools import lru_cache

@lru_cache(maxsize=128, typed=True)
class MyClass:
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 15.7μs -> 4.07μs (286% faster)

    def test_base_class_with_attribute_access_and_method(self):
        """Test class with attribute-accessed base class"""
        source = """import collections

class MyClass(collections.abc.Mapping):
    pass
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 13.1μs -> 3.49μs (276% faster)

    def test_multiple_imports_same_module(self):
        """Test multiple imports from same module"""
        source = """from typing import List, Dict, Optional

class MyClass:
    items: List[str]
    mapping: Dict[str, int]
    value: Optional[str]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 31.7μs -> 6.43μs (393% faster)

    def test_import_before_and_after_class(self):
        """Test correct extraction when imports are interspersed"""
        source = """from typing import List

class MyClass:
    items: List[str]

from os import path
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 16.3μs -> 4.66μs (249% faster)

    def test_class_with_field_call(self):
        """Test class with dataclass field() calls in annotations"""
        source = """from dataclasses import dataclass, field
from typing import List

@dataclass
class MyClass:
    items: List[str] = field(default_factory=list)
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[2]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 22.4μs -> 5.16μs (333% faster)

class TestExtractImportsLargeScale:
    """Test performance and scalability with larger inputs"""

    def test_class_with_many_type_annotations(self):
        """Test class with many type-annotated fields"""
        # Create a class with 50 annotated fields
        imports = "from typing import List, Dict, Optional, Tuple, Set\n\n"
        class_def = "class LargeClass:\n"
        for i in range(50):
            # Rotate through different type annotations
            if i % 5 == 0:
                class_def += f"    field{i}: List[str]\n"
            elif i % 5 == 1:
                class_def += f"    field{i}: Dict[str, int]\n"
            elif i % 5 == 2:
                class_def += f"    field{i}: Optional[str]\n"
            elif i % 5 == 3:
                class_def += f"    field{i}: Tuple[int, str]\n"
            else:
                class_def += f"    field{i}: Set[int]\n"

        source = imports + class_def
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 343μs -> 32.8μs (949% faster)

    def test_class_with_many_decorators(self):
        """Test class with many decorators"""
        decorators = ""
        imports = ""
        for i in range(30):
            decorator_name = f"decorator{i}"
            decorators += f"@{decorator_name}\n"
            imports += f"from module{i % 10} import {decorator_name}\n"

        source = imports + decorators + "class LargeClass:\n    pass\n"
        module_tree = ast.parse(source)
        class_node = module_tree.body[30]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 59.4μs -> 20.4μs (192% faster)
        # Should extract many imports for decorators
        lines = result.split("\n")

    def test_class_with_many_base_classes(self):
        """Test class with many base classes"""
        imports = ""
        bases = []
        for i in range(25):
            base_name = f"Base{i}"
            bases.append(base_name)
            imports += f"from base_module{i % 10} import {base_name}\n"

        base_list = ", ".join(bases)
        source = imports + f"class LargeClass({base_list}):\n    pass\n"
        module_tree = ast.parse(source)
        class_node = module_tree.body[25]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 49.5μs -> 16.6μs (198% faster)
        # Should extract imports for many base classes
        lines = result.split("\n")

    def test_large_module_with_many_imports(self):
        """Test extraction from module with many imports"""
        # Create module with 100 imports, only a few used by target class
        imports = ""
        for i in range(100):
            imports += f"from module{i % 20} import Unused{i}\n"

        # Add the imports we actually need
        imports += "from typing import List\n"
        imports += "from abc import ABC\n"

        source = imports + "class TargetClass(ABC):\n    items: List[str]\n"
        module_tree = ast.parse(source)
        class_node = module_tree.body[102]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 43.9μs -> 29.8μs (47.2% faster)

    def test_very_long_source_file(self):
        """Test with a very long source file"""
        # Create a source file with many lines
        lines = []
        for i in range(200):
            lines.append(f"# Comment line {i}")

        lines.append("from typing import List")
        lines.append("from abc import ABC")
        lines.append("")
        lines.append("class TargetClass(ABC):")
        lines.append("    items: List[str]")

        source = "\n".join(lines)
        module_tree = ast.parse(source)
        class_node = module_tree.body[2]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 25.7μs -> 12.2μs (111% faster)

    def test_complex_nested_annotations_performance(self):
        """Test performance with deeply nested type annotations"""
        # Create annotation like Dict[str, List[Tuple[Optional[int], Set[str]]]]
        source = """from typing import Dict, List, Tuple, Optional, Set

class ComplexClass:
    field1: Dict[str, List[Tuple[Optional[int], Set[str]]]]
    field2: Dict[List[Tuple[Set[Dict[str, Optional[int]]]]], str]
    field3: List[Dict[str, Tuple[Optional[Set[int]], Dict[str, List[str]]]]]
"""
        module_tree = ast.parse(source)
        class_node = module_tree.body[1]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 83.8μs -> 12.1μs (591% faster)

    def test_large_class_with_mixed_features(self):
        """Test class combining many features at scale"""
        imports = "from typing import Dict, List, Optional\n"
        imports += "from abc import ABC\n"
        imports += "from dataclasses import dataclass, field\n"

        decorators = "@dataclass\n"

        bases = "ABC"

        fields = ""
        for i in range(30):
            fields += f"    field{i}: Optional[List[Dict[str, int]]]\n"

        source = imports + decorators + f"class LargeClass({bases}):\n" + fields

        module_tree = ast.parse(source)
        class_node = module_tree.body[3]
        codeflash_output = extract_imports_for_class(module_tree, class_node, source); result = codeflash_output # 406μs -> 41.5μs (880% faster)
# 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-pr1339-2026-02-04T01.09.46 and push.