⚡️ Speed up function extract_imports_for_class by 396% in PR #1166 (skyvern-grace)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1166

If you approve this dependent PR, these changes will be merged into the original PR branch skyvern-grace.

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

---

📄 396% (3.96x) speedup for extract_imports_for_class in codeflash/context/code_context_extractor.py

⏱️ Runtime : 3.23 milliseconds → 650 microseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 396% speedup (3.23ms → 650μs) by replacing an expensive AST traversal with a direct iteration over class body nodes.

Key Optimization

Replaced ast.walk(class_node) with direct class_node.body iteration (line 33):

• Original: Used ast.walk(class_node) which recursively traverses ALL nodes in the class AST (3,785 hits), including method bodies, nested statements, and deeply nested expressions. This accounted for 70.5% of total runtime.
• Optimized: Directly iterates over class_node.body, which contains only the top-level class members (553 hits) - a 7x reduction in nodes visited.

Why This Works

The function only needs to inspect field definitions at the class level to collect type annotation names. Method bodies and nested structures are irrelevant for extracting imports. By iterating only class_node.body:

• We examine just the annotated field assignments (ast.AnnAssign) and field calls needed for import extraction
• We skip irrelevant AST nodes like method definitions, nested statements, and expression details inside methods
• The reduction from 3,785 to 553 node checks directly translates to the observed speedup

Performance Characteristics

Based on the test results, the optimization excels across all scenarios:

• Simple classes: 176-362% speedup (basic imports/decorators)
• Complex nested annotations: 280-586% speedup (Dict[List[Optional[...]]])
• Large-scale scenarios: Up to 1991% speedup for classes with 100+ methods and fields (where the original's deep traversal penalty was most severe)

Impact Assessment

The function is called from get_imported_class_definitions() which extracts class definitions for LLM context during code optimization. This is in a hot path that processes every imported class in the codebase being analyzed. With the 4-5x speedup, code context extraction becomes significantly faster, improving the overall optimization pipeline's responsiveness, especially for large codebases with many dataclass-style classes.

The optimization preserves exact functionality - it still collects all needed import names from base classes, decorators, and type annotations, just by examining the relevant nodes directly rather than walking the entire AST tree.

✅ Correctness verification report:

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

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_code_context_extractor.py::TestExtractImportsForClass.test_extracts_base_class_imports	12.6μs	5.40μs	134%✅
test_code_context_extractor.py::TestExtractImportsForClass.test_extracts_decorator_imports	14.9μs	5.10μs	193%✅
test_code_context_extractor.py::TestExtractImportsForClass.test_extracts_field_function_imports	22.6μs	5.94μs	281%✅
test_code_context_extractor.py::TestExtractImportsForClass.test_extracts_type_annotation_imports	24.7μs	7.22μs	242%✅
test_code_context_extractor.py::TestExtractImportsForClass.test_no_duplicate_imports	23.1μs	6.08μs	281%✅

🌀 Click to see Generated Regression Tests

from __future__ import annotations

# imports
import ast
import textwrap

from codeflash.context.code_context_extractor import extract_imports_for_class

# unit tests


def _get_class_node(module_tree: ast.Module, classname: str) -> ast.ClassDef:
    """Helper to find the class node by name in the module AST."""
    for node in module_tree.body:
        if isinstance(node, ast.ClassDef) and node.name == classname:
            return node
    raise AssertionError(f"Class {classname} not found in AST")


def test_basic_imports_for_simple_dataclass():
    # Basic scenario: class with base class, decorator, and annotated fields using imported names.
    module_source = textwrap.dedent(
        """\
        import abc
        from dataclasses import dataclass, field
        from typing import List

        @dataclass
        class Foo(abc.ABC):
            x: List[int]
            y: int = field(default=0)
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "Foo")

    # Call the function under test
    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 31.6μs -> 6.90μs (358% faster)

    # The function should include the import lines that provide abc, dataclass/field, and List
    lines = result.splitlines()


def test_decorator_call_and_attribute_decorator_detection():
    # Edge scenario: decorator used as a call and as an attribute (module.decorator)
    module_source = textwrap.dedent(
        """\
        import pkg

        @pkg.decorate(arg=1)
        class A:
            pass
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "A")

    # Should detect that the decorator references 'pkg' and therefore return the import for pkg
    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 17.2μs -> 4.34μs (296% faster)


def test_alias_imports_and_attribute_annotations():
    # Edge: import alias (import numpy as np) and attribute annotation (np.ndarray)
    module_source = textwrap.dedent(
        """\
        import numpy as np

        class B:
            arr: np.ndarray
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "B")

    # collect_names_from_annotation will add 'np' from the attribute, matching alias 'np'
    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 15.8μs -> 4.28μs (270% faster)


def test_importfrom_alias_mismatch_behavior():
    # Edge case that demonstrates current behavior: aliasing on import-from causes a mismatch
    module_source = textwrap.dedent(
        """\
        from typing import List as L

        class C:
            a: List[int]
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "C")

    # The function uses alias.asname if present when matching names from ImportFrom.
    # Since annotation referred to 'List' but the import uses alias 'L', the import will NOT be detected.
    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 16.6μs -> 4.11μs (304% faster)


def test_importfrom_alias_match_when_using_alias_name():
    # Complementary check: if the annotation uses the alias name, it should find the import.
    module_source = textwrap.dedent(
        """\
        from typing import List as L

        class C2:
            a: L[int]
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "C2")

    # Now annotation uses alias 'L' so the import-from line should be returned
    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 17.3μs -> 4.44μs (290% faster)


def test_union_and_tuple_and_subscript_annotations():
    # Complex annotations: Union via | operator, subscripted types, and tuple in subscripts.
    module_source = textwrap.dedent(
        """\
        from modA import A
        from modB import B, C

        class D:
            a: A | B
            b: tuple[C, B]
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "D")

    # The union (A | B) should add A and B; tuple[C, B] should add 'tuple', C, and B
    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 28.2μs -> 6.86μs (312% faster)
    lines = result.splitlines()


def test_field_call_detection_and_deduplication():
    # Dataclass field() used multiple times but import must appear only once in output
    module_source = textwrap.dedent(
        """\
        from dataclasses import dataclass, field

        @dataclass
        class E:
            a: int = field(default=1)
            b: int = field(default=2)
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "E")

    # The single "from dataclasses..." line should be present only once even though field() is used twice
    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 27.2μs -> 4.32μs (530% faster)
    lines = result.splitlines()


def test_empty_class_requires_no_imports():
    # Edge: class with no bases, no decorators, and no annotations should return an empty string
    module_source = textwrap.dedent(
        """\
        import something

        class Empty:
            pass
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "Empty")

    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 9.23μs -> 2.87μs (222% faster)


def test_large_scale_many_imports_and_many_annotations():
    # Large scale scenario: create many import lines and a class that references a subset of them.
    # We generate 200 import lines and reference 100 of them in the class body.
    num_imports = 200
    num_referenced = 100  # number of different imports the class will actually need
    import_lines = [f"import mod{i}" for i in range(num_imports)]
    # Create annotations referencing mod0.Type0, mod2.Type2, ..., mod198.Type198 for even indices up to num_referenced*2
    referenced_indices = list(range(0, num_referenced * 2, 2))  # even numbers: 0,2,4,...
    class_body_lines = []
    for idx in referenced_indices:
        # Each annotated attribute uses an attribute-style annotation so collect_names_from_annotation extracts module name
        class_body_lines.append(f"    f{idx}: mod{idx}.Type{idx}")
    module_source = "\n".join(import_lines) + "\n\nclass Big:\n" + "\n".join(class_body_lines) + "\n"
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "Big")

    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 595μs -> 164μs (261% faster)
    lines = [line for line in result.splitlines() if line.strip()]


def test_annotation_tuple_in_parentheses_is_collected():
    # Edge: annotations written using parentheses that produce an AST.Tuple in the annotation
    module_source = textwrap.dedent(
        """\
        from pkg1 import X, Y

        class T:
            a: (X, Y)
        """
    )
    module_tree = ast.parse(module_source)
    class_node = _get_class_node(module_tree, "T")

    # The tuple annotation should lead to both X and Y being collected and the import returned
    codeflash_output = extract_imports_for_class(module_tree, class_node, module_source)
    result = codeflash_output  # 17.7μs -> 4.61μs (284% 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

from codeflash.context.code_context_extractor import extract_imports_for_class


def test_basic_single_base_class_import():
    """Test extraction of import for a simple base class."""
    source = """import BaseClass

class MyClass(BaseClass):
    pass
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 15.2μs -> 5.51μs (176% faster)


def test_basic_decorator_import():
    """Test extraction of import for a decorator."""
    source = """from dataclasses import dataclass

@dataclass
class MyClass:
    pass
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 13.0μs -> 4.58μs (185% faster)


def test_basic_type_annotation_import():
    """Test extraction of import for type annotation in class field."""
    source = """from typing import List

class MyClass:
    items: List[str]
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 18.2μs -> 5.77μs (215% faster)


def test_basic_no_imports_needed():
    """Test class with only builtin base (object) requires no imports."""
    source = """class MyClass(object):
    pass
"""
    module = ast.parse(source)
    class_node = module.body[0]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 10.6μs -> 2.77μs (284% faster)


def test_edge_empty_class():
    """Test extraction from empty class with no bases or decorators."""
    source = """class EmptyClass:
    pass
"""
    module = ast.parse(source)
    class_node = module.body[0]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 12.1μs -> 3.17μs (281% faster)


def test_edge_nested_type_annotations():
    """Test extraction from complex nested type annotations."""
    source = """from typing import List, Dict, Optional

class MyClass:
    data: Dict[str, List[Optional[int]]]
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 27.7μs -> 7.29μs (280% faster)


def test_edge_attribute_base_class():
    """Test base class accessed via attribute (e.g., abc.ABC)."""
    source = """import abc

class MyClass(abc.ABC):
    pass
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 13.7μs -> 4.43μs (210% faster)


def test_edge_decorator_with_arguments():
    """Test decorator function called with arguments."""
    source = """from dataclasses import dataclass, field

@dataclass(frozen=True)
class MyClass:
    value: int = field(default=0)
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 23.4μs -> 5.07μs (362% faster)


def test_edge_decorator_with_attribute_access():
    """Test decorator accessed via module attribute."""
    source = """import functools

@functools.lru_cache
class MyClass:
    pass
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 12.6μs -> 3.38μs (274% faster)


def test_edge_multiple_base_classes():
    """Test class with multiple base classes."""
    source = """from abc import ABC
from collections.abc import Iterable

class MyClass(ABC, Iterable):
    pass
"""
    module = ast.parse(source)
    class_node = module.body[2]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 13.9μs -> 4.83μs (187% faster)


def test_edge_import_with_alias():
    """Test import with alias (as)."""
    source = """from dataclasses import dataclass as dc

@dc
class MyClass:
    pass
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 11.7μs -> 3.88μs (203% faster)


def test_edge_import_with_builtin_names():
    """Test that builtin names don't appear in extracted imports."""
    source = """class MyClass:
    data: list
    value: dict
"""
    module = ast.parse(source)
    class_node = module.body[0]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 16.3μs -> 3.38μs (382% faster)


def test_edge_union_type_with_pipe_syntax():
    """Test Union type using pipe syntax (Python 3.10+)."""
    source = """class MyClass:
    value: int | str
"""
    module = ast.parse(source)
    class_node = module.body[0]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 15.8μs -> 3.69μs (329% faster)


def test_edge_generic_type_annotation():
    """Test generic type annotation extraction."""
    source = """from typing import TypeVar, Generic

T = TypeVar('T')

class MyClass(Generic[T]):
    value: T
"""
    module = ast.parse(source)
    class_node = module.body[2]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 18.2μs -> 4.67μs (289% faster)


def test_edge_no_duplicate_imports():
    """Test that duplicate imports are not included twice."""
    source = """from typing import List

class MyClass:
    items1: List[int]
    items2: List[str]
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 24.0μs -> 5.65μs (325% faster)
    # Count occurrences - should only appear once
    count = result.count("from typing import List")


def test_edge_unused_imports_not_included():
    """Test that unused imports are not extracted."""
    source = """import os
from dataclasses import dataclass

@dataclass
class MyClass:
    pass
"""
    module = ast.parse(source)
    class_node = module.body[2]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 12.4μs -> 4.60μs (169% faster)


def test_edge_star_import():
    """Test behavior with star imports."""
    source = """from typing import *

class MyClass:
    items: List
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 13.5μs -> 3.50μs (286% faster)


def test_edge_relative_import():
    """Test handling of relative imports."""
    source = """from .module import BaseClass

class MyClass(BaseClass):
    pass
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 11.3μs -> 3.72μs (203% faster)


def test_edge_import_multiple_names_one_used():
    """Test import statement with multiple names when only one is used."""
    source = """from dataclasses import dataclass, field, FrozenInstanceError

@dataclass
class MyClass:
    pass
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 11.9μs -> 3.85μs (209% faster)


def test_large_many_base_classes():
    """Test class inheriting from many base classes."""
    # Create a source with many base classes
    base_names = [f"Base{i}" for i in range(50)]
    imports = "\n".join([f"from module{i} import Base{i}" for i in range(50)])
    bases_str = ", ".join(base_names)

    source = f"""{imports}

class MyClass({bases_str}):
    pass
"""
    module = ast.parse(source)
    class_node = module.body[50]  # After 50 imports
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 99.2μs -> 33.9μs (193% faster)

    # All base classes should be represented in the result
    for i in range(50):
        pass


def test_large_many_decorators():
    """Test class with many decorators."""
    decorators = "\n".join([f"@decorator{i}" for i in range(100)])
    imports = "\n".join([f"from module{i} import decorator{i}" for i in range(100)])

    source = f"""{imports}

{decorators}
class MyClass:
    pass
"""
    module = ast.parse(source)
    class_node = module.body[100]  # After 100 imports
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 184μs -> 61.2μs (202% faster)

    # All decorators should be represented
    for i in range(100):
        pass


def test_large_many_type_annotations():
    """Test class with many type-annotated fields."""
    imports = "from typing import List\n" + "\n".join([f"import module{i}" for i in range(100)])
    annotations = "\n".join([f"    field{i}: List" for i in range(100)])

    source = f"""{imports}

class MyClass:
{annotations}
"""
    module = ast.parse(source)
    class_node = module.body[101]  # After imports
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 390μs -> 63.2μs (518% faster)


def test_large_complex_nested_annotations():
    """Test class with complex nested type annotations."""
    source = """from typing import Dict, List, Optional, Union, Tuple

class MyClass:
    data1: Dict[str, List[int]]
    data2: Optional[Union[int, str]]
    data3: Tuple[int, str, float]
    data4: List[Dict[str, Optional[int]]]
    data5: Union[Dict[str, int], List[str]]
    data6: Optional[Tuple[List[int], Dict[str, str]]]
    data7: Dict[str, Union[int, str]]
    data8: List[Optional[Union[int, str]]]
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 125μs -> 18.3μs (586% faster)


def test_large_many_imports_in_module():
    """Test extracting from module with many imports."""
    # Create a module with many imports, only a few used
    imports_list = []
    for i in range(200):
        imports_list.append(f"import module{i}")

    # Only use a few
    source = (
        "\n".join(imports_list)
        + """

class MyClass(module5, module50, module150):
    pass
"""
    )
    module = ast.parse(source)
    class_node = module.body[200]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 77.9μs -> 64.4μs (20.9% faster)


def test_large_mixed_complex_scenario():
    """Test large scenario with decorators, bases, and annotations."""
    imports = """from dataclasses import dataclass, field
from typing import List, Dict, Optional
import abc
from collections.abc import Iterable
"""

    decorators = "@dataclass\n@abc.abstractmethod" if False else "@dataclass"

    source = f"""{imports}

{decorators}
class MyClass(abc.ABC, Iterable):
    items: List[Dict[str, Optional[int]]]
    metadata: Dict[str, str] = field(default_factory=dict)
"""
    module = ast.parse(source)
    class_node = module.body[4]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 49.1μs -> 10.7μs (360% faster)


def test_large_performance_with_deep_ast_walk():
    """Test performance with deeply nested class structure."""
    # Create a class with many nested method definitions and assignments
    methods = "\n".join([f"    def method{i}(self) -> int: return {i}" for i in range(100)])
    annotations = "\n".join([f"    field{i}: int" for i in range(100)])

    source = f"""from typing import Optional

class MyClass:
{annotations}
{methods}
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 1.03ms -> 49.2μs (1991% faster)


def test_edge_callable_type_annotation():
    """Test Callable type annotation extraction."""
    source = """from typing import Callable

class MyClass:
    callback: Callable[[int, str], bool]
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 24.1μs -> 6.01μs (302% faster)


def test_edge_module_level_constant_in_annotation():
    """Test extraction when annotation references module-level constant."""
    source = """from typing import Literal

MyValue = "constant"

class MyClass:
    status: Literal["active", "inactive"]
"""
    module = ast.parse(source)
    class_node = module.body[2]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 21.0μs -> 6.20μs (238% faster)


def test_edge_forward_reference_in_annotation():
    """Test handling of string annotations (forward references)."""
    source = """from typing import Optional

class MyClass:
    ref: Optional['MyClass']
"""
    module = ast.parse(source)
    class_node = module.body[1]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 16.9μs -> 5.01μs (237% faster)


def test_edge_multiple_decorator_calls():
    """Test multiple decorator calls on same class."""
    source = """from functools import wraps
from dataclasses import dataclass

@wraps
@dataclass
class MyClass:
    pass
"""
    module = ast.parse(source)
    class_node = module.body[2]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 14.3μs -> 4.80μs (197% faster)


def test_edge_field_default_factory():
    """Test field() with default_factory parameter."""
    source = """from dataclasses import dataclass, field
from typing import List

@dataclass
class MyClass:
    items: List[str] = field(default_factory=list)
"""
    module = ast.parse(source)
    class_node = module.body[2]
    codeflash_output = extract_imports_for_class(module, class_node, source)
    result = codeflash_output  # 24.5μs -> 5.68μs (331% 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-pr1166-2026-01-24T09.06.26 and push.