⚡️ Speed up method InitDecorator.visit_ClassDef by 91% in PR #363 (part-1-windows-fixes)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #363

If you approve this dependent PR, these changes will be merged into the original PR branch part-1-windows-fixes.

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

---

📄 91% (0.91x) speedup for InitDecorator.visit_ClassDef in codeflash/verification/instrument_codeflash_capture.py

⏱️ Runtime : 548 microseconds → 287 microseconds (best of 104 runs)

📝 Explanation and details

Here is an optimized version of your program with reduced code execution overhead, memory allocation, and efficient logic. All return values remain exactly the same, and function signatures are unchanged. Comments are preserved as required.

Key Optimizations:

• Convert repeated property accesses/calls to local variables where possible.
• Reuse the same decorator AST node object where safe (AST is not mutated elsewhere).
• Pre-check presence of target class and short-circuit early.
• In the loop, short-circuit when __init__ is found to avoid unnecessary iterations.

Summary of improvements:

• Reused AST decorator components, minimizing repeated object creation.
• Early exit on target class and __init__ findings.
• Less variable assignment and number of iterations.
• All core logic and comments are preserved.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ast
import sys
from pathlib import Path

# imports
import pytest  # used for our unit tests
from codeflash.verification.instrument_codeflash_capture import InitDecorator

# -------------------------
# Helper functions for tests
# -------------------------

def get_classdef_node(source: str, class_name: str) -> ast.ClassDef:
    """Parse source and return the ast.ClassDef node with the given name."""
    mod = ast.parse(source)
    for node in mod.body:
        if isinstance(node, ast.ClassDef) and node.name == class_name:
            return node
    raise ValueError(f"Class {class_name} not found in AST.")

def get_init_func(node: ast.ClassDef) -> ast.FunctionDef | None:
    """Return the __init__ function node from a classdef, or None if not present."""
    for item in node.body:
        if isinstance(item, ast.FunctionDef) and item.name == "__init__":
            return item
    return None

def get_codeflash_capture_decorator(func: ast.FunctionDef) -> ast.Call | None:
    """Return the codeflash_capture decorator node if present, else None."""
    for deco in func.decorator_list:
        if (
            isinstance(deco, ast.Call)
            and isinstance(deco.func, ast.Name)
            and deco.func.id == "codeflash_capture"
        ):
            return deco
    return None

def decorator_args_match(deco: ast.Call, class_name: str, tmp_dir_path: str, tests_root: str, is_fto: bool) -> bool:
    """Check that the decorator's arguments match the expected values."""
    kw = {k.arg: k.value for k in deco.keywords}
    return (
        isinstance(kw["function_name"], ast.Constant)
        and kw["function_name"].value == f"{class_name}.__init__"
        and isinstance(kw["tmp_dir_path"], ast.Constant)
        and kw["tmp_dir_path"].value == tmp_dir_path
        and isinstance(kw["tests_root"], ast.Constant)
        and kw["tests_root"].value == tests_root
        and isinstance(kw["is_fto"], ast.Constant)
        and kw["is_fto"].value == is_fto
    )

# -------------------------
# Unit tests
# -------------------------

# Basic Test Cases

def test_adds_decorator_to_existing_init():
    """Test that the decorator is added to a class with a plain __init__."""
    source = """
class Foo:
    def __init__(self, x):
        self.x = x
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    class_name = "Foo"
    node = get_classdef_node(source, class_name)
    decorator = InitDecorator({class_name}, "fto", tmp_dir, tests_root, is_fto=True)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    deco = get_codeflash_capture_decorator(init_func)

def test_does_not_modify_other_classes():
    """Test that classes not in target_classes are left untouched."""
    source = """
class Bar:
    def __init__(self, y):
        self.y = y
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "Bar")
    decorator = InitDecorator({"Foo"}, "fto", tmp_dir, tests_root, is_fto=False)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    deco = get_codeflash_capture_decorator(init_func)

def test_adds_init_if_missing():
    """Test that __init__ is created if missing, and decorated."""
    source = """
class Baz:
    pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "Baz")
    decorator = InitDecorator({"Baz"}, "fto", tmp_dir, tests_root, is_fto=False)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    deco = get_codeflash_capture_decorator(init_func)

def test_existing_init_with_other_decorators():
    """Test that codeflash_capture decorator is inserted at the front if other decorators exist."""
    source = """
class Quux:
    @staticmethod
    def __init__(self):
        pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "Quux")
    decorator = InitDecorator({"Quux"}, "fto", tmp_dir, tests_root, is_fto=True)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    # codeflash_capture should be first
    first_deco = init_func.decorator_list[0]

def test_existing_init_with_self_not_first_arg():
    """Test that __init__ with wrong first arg is not decorated."""
    source = """
class WrongInit:
    def __init__(notself):
        pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "WrongInit")
    decorator = InitDecorator({"WrongInit"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    deco = get_codeflash_capture_decorator(init_func)

# Edge Test Cases

def test_init_with_no_args():
    """Test that __init__ with no args is not decorated (invalid signature)."""
    source = """
class EmptyInit:
    def __init__():
        pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "EmptyInit")
    decorator = InitDecorator({"EmptyInit"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    deco = get_codeflash_capture_decorator(init_func)

def test_multiple_classes_some_targeted():
    """Test that only targeted classes are modified in a multi-class module."""
    source = """
class A:
    def __init__(self): pass
class B:
    pass
class C:
    def __init__(self): pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    # Parse the module and get all class nodes
    mod = ast.parse(source)
    class_nodes = {node.name: node for node in mod.body if isinstance(node, ast.ClassDef)}
    decorator = InitDecorator({"A", "C"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(class_nodes["A"]); result_A = codeflash_output
    codeflash_output = decorator.visit_ClassDef(class_nodes["B"]); result_B = codeflash_output
    codeflash_output = decorator.visit_ClassDef(class_nodes["C"]); result_C = codeflash_output
    # A and C should be decorated
    for name, result in [("A", result_A), ("C", result_C)]:
        init_func = get_init_func(result)
        deco = get_codeflash_capture_decorator(init_func)
    # B should not be decorated
    init_func_B = get_init_func(result_B)
    if init_func_B:
        deco_B = get_codeflash_capture_decorator(init_func_B)

def test_existing_codeflash_capture_not_duplicated():
    """Test that an existing codeflash_capture decorator is not duplicated."""
    source = """
class AlreadyDecorated:
    @codeflash_capture(function_name="AlreadyDecorated.__init__", tmp_dir_path="/tmp", tests_root="/tests", is_fto=False)
    def __init__(self):
        pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "AlreadyDecorated")
    decorator = InitDecorator({"AlreadyDecorated"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    # There should be only one codeflash_capture decorator
    deco_count = sum(
        1 for d in init_func.decorator_list
        if isinstance(d, ast.Call) and isinstance(d.func, ast.Name) and d.func.id == "codeflash_capture"
    )

def test_class_with_other_methods_but_no_init():
    """Test that __init__ is inserted before other methods."""
    source = """
class InsertInit:
    def foo(self): pass
    def bar(self): pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "InsertInit")
    decorator = InitDecorator({"InsertInit"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    deco = get_codeflash_capture_decorator(result.body[0])

def test_class_with_init_and_classmethod():
    """Test that only the instance __init__ is decorated, not classmethods."""
    source = """
class ClassWithClassmethod:
    @classmethod
    def __init__(cls):
        pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "ClassWithClassmethod")
    decorator = InitDecorator({"ClassWithClassmethod"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    # Should not decorate classmethod __init__
    deco = get_codeflash_capture_decorator(init_func)

def test_class_with_init_and_staticmethod():
    """Test that only the instance __init__ is decorated, not staticmethods."""
    source = """
class ClassWithStaticmethod:
    @staticmethod
    def __init__():
        pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "ClassWithStaticmethod")
    decorator = InitDecorator({"ClassWithStaticmethod"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    # Should not decorate staticmethod __init__
    deco = get_codeflash_capture_decorator(init_func)

def test_class_with_init_with_varargs_and_kwargs():
    """Test that __init__ with *args and **kwargs is decorated."""
    source = """
class VarArgsInit:
    def __init__(self, *args, **kwargs):
        pass
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "VarArgsInit")
    decorator = InitDecorator({"VarArgsInit"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    deco = get_codeflash_capture_decorator(init_func)

# Large Scale Test Cases


def test_large_class_body_with_existing_init():
    """Test that decorator is added to __init__ in a class with a large body."""
    NUM_METHODS = 500
    methods = "\n".join([f"    def method{i}(self): pass" for i in range(NUM_METHODS)])
    source = f"""
class BigClass:
    def __init__(self): pass
{methods}
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "BigClass")
    decorator = InitDecorator({"BigClass"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    deco = get_codeflash_capture_decorator(init_func)

def test_large_class_body_without_init():
    """Test that __init__ is inserted and decorated in a large class without __init__."""
    NUM_METHODS = 500
    methods = "\n".join([f"    def method{i}(self): pass" for i in range(NUM_METHODS)])
    source = f"""
class BigClassNoInit:
{methods}
"""
    tmp_dir = "/tmp"
    tests_root = Path("/tests")
    node = get_classdef_node(source, "BigClassNoInit")
    decorator = InitDecorator({"BigClassNoInit"}, "fto", tmp_dir, tests_root)
    codeflash_output = decorator.visit_ClassDef(node); result = codeflash_output
    init_func = get_init_func(result)
    deco = get_codeflash_capture_decorator(init_func)



from __future__ import annotations

import ast
import sys
from pathlib import Path

# imports
import pytest  # used for our unit tests
from codeflash.verification.instrument_codeflash_capture import InitDecorator

# ----------------- UNIT TESTS -----------------

# Helper function to extract class AST node from code
def get_class_node(source: str, class_name: str) -> ast.ClassDef:
    """Parse source and return the ast.ClassDef node for class_name."""
    tree = ast.parse(source)
    for node in tree.body:
        if isinstance(node, ast.ClassDef) and node.name == class_name:
            return node
    raise ValueError(f"Class {class_name} not found in AST")


# Helper function to get __init__ function from class node
def get_init_func(class_node: ast.ClassDef) -> ast.FunctionDef | None:
    for item in class_node.body:
        if isinstance(item, ast.FunctionDef) and item.name == "__init__":
            return item
    return None

# Helper function to check if codeflash_capture is present as decorator
def has_codeflash_capture_decorator(init_func: ast.FunctionDef) -> bool:
    for d in init_func.decorator_list:
        if isinstance(d, ast.Call) and isinstance(d.func, ast.Name) and d.func.id == "codeflash_capture":
            return True
    return False

# Helper function to get the codeflash_capture decorator call
def get_codeflash_capture_decorator(init_func: ast.FunctionDef):
    for d in init_func.decorator_list:
        if isinstance(d, ast.Call) and isinstance(d.func, ast.Name) and d.func.id == "codeflash_capture":
            return d
    return None

# Helper function to check for super().__init__(*args, **kwargs) in __init__ body
def has_super_init_call(init_func: ast.FunctionDef) -> bool:
    for stmt in init_func.body:
        if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
            call = stmt.value
            if (isinstance(call.func, ast.Attribute) and
                isinstance(call.func.value, ast.Call) and
                isinstance(call.func.value.func, ast.Name) and
                call.func.value.func.id == "super" and
                call.func.attr == "__init__"):
                # Check for *args and **kwargs
                found_args = any(isinstance(arg, ast.Starred) and isinstance(arg.value, ast.Name) and arg.value.id == "args"
                                 for arg in call.args)
                found_kwargs = any(isinstance(kw.value, ast.Name) and kw.arg is None and kw.value.id == "kwargs"
                                   for kw in call.keywords)
                return found_args and found_kwargs
    return False

# ----------------- BASIC TEST CASES -----------------

def test_adds_decorator_to_existing_init():
    """Test: Adds decorator to a class with an existing __init__."""
    src = """
class Foo:
    def __init__(self, x):
        self.x = x
    def bar(self):
        pass
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
        is_fto=True,
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)
    # Check that decorator arguments are correct
    deco = get_codeflash_capture_decorator(init_func)
    kw = {k.arg: k.value for k in deco.keywords}

def test_does_not_add_decorator_to_non_target_class():
    """Test: Does not modify class if not in target_classes."""
    src = """
class Bar:
    def __init__(self, y):
        self.y = y
"""
    class_node = get_class_node(src, "Bar")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_adds_init_if_missing():
    """Test: Adds __init__ with decorator if missing."""
    src = """
class Baz:
    def hello(self): pass
"""
    class_node = get_class_node(src, "Baz")
    decorator = InitDecorator(
        target_classes={"Baz"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_decorator_not_duplicated():
    """Test: Does not duplicate codeflash_capture if already present."""
    src = """
class Foo:
    @codeflash_capture(function_name="Foo.__init__", tmp_dir_path="/tmp", tests_root="/tests", is_fto=True)
    def __init__(self, x):
        self.x = x
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
        is_fto=True,
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)
    # Should only have one codeflash_capture decorator
    count = sum(
        1 for d in init_func.decorator_list
        if isinstance(d, ast.Call) and isinstance(d.func, ast.Name) and d.func.id == "codeflash_capture"
    )

# ----------------- EDGE TEST CASES -----------------

def test_init_with_no_args():
    """Test: Handles __init__ with no arguments except self."""
    src = """
class Foo:
    def __init__(self):
        pass
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_init_with_varargs_and_kwargs():
    """Test: Handles __init__ with *args and **kwargs."""
    src = """
class Foo:
    def __init__(self, *args, **kwargs):
        pass
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_init_with_positional_and_keyword_only_args():
    """Test: Handles __init__ with positional and keyword-only args."""
    src = """
class Foo:
    def __init__(self, a, *, b=None):
        pass
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_class_with_multiple_methods_and_no_init():
    """Test: Adds __init__ if missing, even with multiple other methods."""
    src = """
class Foo:
    def a(self): pass
    def b(self): pass
    def c(self): pass
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_class_with_init_but_no_self():
    """Test: Does not add decorator if __init__ is missing self (should not match)."""
    src = """
class Foo:
    def __init__(x):
        pass
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)
    # Should add a new __init__ with decorator, since the existing one is not valid
    # (according to the matching logic: must have self as first arg)
    # So, should have two __init__ methods
    count = sum(1 for item in new_node.body if isinstance(item, ast.FunctionDef) and item.name == "__init__")
    # The first should be the one inserted
    inserted_init = new_node.body[0]

def test_class_with_no_body():
    """Test: Handles class with no body (edge case)."""
    src = """
class Foo:
    pass
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_class_with_docstring_only():
    """Test: Handles class with only a docstring."""
    src = '''
class Foo:
    """This is a docstring."""
'''
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_class_with_nested_class():
    """Test: Only modifies the outer class, not nested classes."""
    src = """
class Foo:
    def __init__(self): pass
    class Bar:
        def __init__(self): pass
"""
    class_node = get_class_node(src, "Foo")
    decorator = InitDecorator(
        target_classes={"Foo"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    # Only the outer class's __init__ should be decorated
    outer_init = get_init_func(new_node)
    # The nested class should not be touched
    nested_class = None
    for item in new_node.body:
        if isinstance(item, ast.ClassDef) and item.name == "Bar":
            nested_class = item
            break
    nested_init = get_init_func(nested_class)

# ----------------- LARGE SCALE TEST CASES -----------------

def test_many_classes_only_target_modified():
    """Test: In a module with many classes, only target class is modified."""
    src = "\n".join(
        f"class C{i}:\n    def __init__(self): pass\n    def foo(self): pass\n"
        for i in range(100)
    )
    class_nodes = []
    tree = ast.parse(src)
    for node in tree.body:
        if isinstance(node, ast.ClassDef):
            class_nodes.append(node)
    # Only target class is C42
    decorator = InitDecorator(
        target_classes={"C42"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    # Apply to all classes
    new_nodes = [decorator.visit_ClassDef(node) for node in class_nodes]
    for i, node in enumerate(new_nodes):
        init_func = get_init_func(node)
        if node.name == "C42":
            pass
        else:
            pass

def test_large_class_with_many_methods():
    """Test: Large class with 500 methods, __init__ is decorated."""
    methods = "\n".join(
        f"    def m{i}(self): pass" for i in range(500)
    )
    src = f"""
class BigClass:
    def __init__(self, x): self.x = x
{methods}
"""
    class_node = get_class_node(src, "BigClass")
    decorator = InitDecorator(
        target_classes={"BigClass"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_large_class_missing_init():
    """Test: Large class with 500 methods, no __init__, adds decorated __init__."""
    methods = "\n".join(
        f"    def m{i}(self): pass" for i in range(500)
    )
    src = f"""
class BigClass:
{methods}
"""
    class_node = get_class_node(src, "BigClass")
    decorator = InitDecorator(
        target_classes={"BigClass"},
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    codeflash_output = decorator.visit_ClassDef(class_node); new_node = codeflash_output
    init_func = get_init_func(new_node)

def test_many_target_classes():
    """Test: Many target classes in one go."""
    src = "\n".join(
        f"class C{i}:\n    def __init__(self): pass\n"
        for i in range(100)
    )
    class_nodes = []
    tree = ast.parse(src)
    for node in tree.body:
        if isinstance(node, ast.ClassDef):
            class_nodes.append(node)
    targets = {f"C{i}" for i in range(0, 100, 10)}  # C0, C10, ..., C90
    decorator = InitDecorator(
        target_classes=targets,
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    new_nodes = [decorator.visit_ClassDef(node) for node in class_nodes]
    for node in new_nodes:
        init_func = get_init_func(node)
        if node.name in targets:
            pass
        else:
            pass

def test_performance_on_large_input():
    """Test: Performance does not degrade on 999 classes."""
    src = "\n".join(
        f"class C{i}:\n    def __init__(self): pass\n"
        for i in range(999)
    )
    class_nodes = []
    tree = ast.parse(src)
    for node in tree.body:
        if isinstance(node, ast.ClassDef):
            class_nodes.append(node)
    targets = {f"C{i}" for i in range(0, 999, 50)}  # every 50th class is a target
    decorator = InitDecorator(
        target_classes=targets,
        fto_name="fto",
        tmp_dir_path="/tmp",
        tests_root=Path("/tests"),
    )
    # Apply to all classes
    new_nodes = [decorator.visit_ClassDef(node) for node in class_nodes]
    for node in new_nodes:
        init_func = get_init_func(node)
        if node.name in targets:
            pass
        else:
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from ast import ClassDef
from codeflash.verification.instrument_codeflash_capture import InitDecorator
from pathlib import Path
import pytest

def test_InitDecorator_visit_ClassDef():
    with pytest.raises(AttributeError, match="'ClassDef'\\ object\\ has\\ no\\ attribute\\ 'name'"):
        InitDecorator.visit_ClassDef(InitDecorator({''}, '', '', Path(), is_fto=0), ClassDef())

To edit these changes git checkout codeflash/optimize-pr363-2025-06-22T22.59.36 and push.