[FIX] Handle new added classes

codeflash/code_utils/code_extractor.py

@@ -72,6 +72,33 @@ def visit_Assign(self, node: cst.Assign) -> Optional[bool]:
         return True
 
 
+def find_insertion_index_after_imports(node: cst.Module) -> int:
+    """Find the position of the last import statement in the top-level of the module."""
+    insert_index = 0
+    for i, stmt in enumerate(node.body):
+        is_top_level_import = isinstance(stmt, cst.SimpleStatementLine) and any(
+            isinstance(child, (cst.Import, cst.ImportFrom)) for child in stmt.body
+        )
+
+        is_conditional_import = isinstance(stmt, cst.If) and all(
+            isinstance(inner, cst.SimpleStatementLine)
+            and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body)
+            for inner in stmt.body.body
+        )
+
+        if is_top_level_import or is_conditional_import:
+            insert_index = i + 1
+
+        # Stop scanning once we reach a class or function definition.
+        # Imports are supposed to be at the top of the file, but they can technically appear anywhere, even at the bottom of the file.
+        # Without this check, a stray import later in the file
+        # would incorrectly shift our insertion index below actual code definitions.
+        if isinstance(stmt, (cst.ClassDef, cst.FunctionDef)):
+            break
+
+    return insert_index
+
+
 class GlobalAssignmentTransformer(cst.CSTTransformer):
     """Transforms global assignments in the original file with those from the new file."""
 
@@ -122,32 +149,6 @@ def leave_Assign(self, original_node: cst.Assign, updated_node: cst.Assign) -> c
 
         return updated_node
 
-    def _find_insertion_index(self, updated_node: cst.Module) -> int:
-        """Find the position of the last import statement in the top-level of the module."""
-        insert_index = 0
-        for i, stmt in enumerate(updated_node.body):
-            is_top_level_import = isinstance(stmt, cst.SimpleStatementLine) and any(
-                isinstance(child, (cst.Import, cst.ImportFrom)) for child in stmt.body
-            )
-
-            is_conditional_import = isinstance(stmt, cst.If) and all(
-                isinstance(inner, cst.SimpleStatementLine)
-                and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body)
-                for inner in stmt.body.body
-            )
-
-            if is_top_level_import or is_conditional_import:
-                insert_index = i + 1
-
-            # Stop scanning once we reach a class or function definition.
-            # Imports are supposed to be at the top of the file, but they can technically appear anywhere, even at the bottom of the file.
-            # Without this check, a stray import later in the file
-            # would incorrectly shift our insertion index below actual code definitions.
-            if isinstance(stmt, (cst.ClassDef, cst.FunctionDef)):
-                break
-
-        return insert_index
-
     def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
         # Add any new assignments that weren't in the original file
         new_statements = list(updated_node.body)
@@ -161,7 +162,7 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
 
         if assignments_to_append:
             # after last top-level imports
-            insert_index = self._find_insertion_index(updated_node)
+            insert_index = find_insertion_index_after_imports(updated_node)
 
             assignment_lines = [
                 cst.SimpleStatementLine([assignment], leading_lines=[cst.EmptyLine()])

codeflash/code_utils/code_replacer.py

@@ -3,13 +3,18 @@
 import ast
 from collections import defaultdict
 from functools import lru_cache
+from itertools import chain
 from typing import TYPE_CHECKING, Optional, TypeVar
 
 import libcst as cst
 from libcst.metadata import PositionProvider
 
 from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.code_extractor import add_global_assignments, add_needed_imports_from_module
+from codeflash.code_utils.code_extractor import (
+    add_global_assignments,
+    add_needed_imports_from_module,
+    find_insertion_index_after_imports,
+)
 from codeflash.code_utils.config_parser import find_conftest_files
 from codeflash.code_utils.formatter import sort_imports
 from codeflash.code_utils.line_profile_utils import ImportAdder
@@ -249,6 +254,7 @@ def __init__(
         ] = {}  # keys are (class_name, function_name)
         self.new_functions: list[cst.FunctionDef] = []
         self.new_class_functions: dict[str, list[cst.FunctionDef]] = defaultdict(list)
+        self.new_classes: list[cst.ClassDef] = []
         self.current_class = None
         self.modified_init_functions: dict[str, cst.FunctionDef] = {}
 
@@ -271,6 +277,10 @@ def visit_ClassDef(self, node: cst.ClassDef) -> bool:
         self.current_class = node.name.value
 
         parents = (FunctionParent(name=node.name.value, type="ClassDef"),)
+
+        if (node.name.value, ()) not in self.preexisting_objects:
+            self.new_classes.append(node)
+
         for child_node in node.body.body:
             if (
                 self.preexisting_objects
@@ -290,13 +300,15 @@ class OptimFunctionReplacer(cst.CSTTransformer):
     def __init__(
         self,
         modified_functions: Optional[dict[tuple[str | None, str], cst.FunctionDef]] = None,
+        new_classes: Optional[list[cst.ClassDef]] = None,
         new_functions: Optional[list[cst.FunctionDef]] = None,
         new_class_functions: Optional[dict[str, list[cst.FunctionDef]]] = None,
         modified_init_functions: Optional[dict[str, cst.FunctionDef]] = None,
     ) -> None:
         super().__init__()
         self.modified_functions = modified_functions if modified_functions is not None else {}
         self.new_functions = new_functions if new_functions is not None else []
+        self.new_classes = new_classes if new_classes is not None else []
         self.new_class_functions = new_class_functions if new_class_functions is not None else defaultdict(list)
         self.modified_init_functions: dict[str, cst.FunctionDef] = (
             modified_init_functions if modified_init_functions is not None else {}
@@ -335,19 +347,33 @@ def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef
     def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
         node = updated_node
         max_function_index = None
-        class_index = None
+        max_class_index = None
         for index, _node in enumerate(node.body):
             if isinstance(_node, cst.FunctionDef):
                 max_function_index = index
             if isinstance(_node, cst.ClassDef):
-                class_index = index
+                max_class_index = index
+
+        if self.new_classes:
+            existing_class_names = {_node.name.value for _node in node.body if isinstance(_node, cst.ClassDef)}
+
+            unique_classes = [
+                new_class for new_class in self.new_classes if new_class.name.value not in existing_class_names
+            ]
+            if unique_classes:
+                new_classes_insertion_idx = max_class_index or find_insertion_index_after_imports(node)
+                new_body = list(
+                    chain(node.body[:new_classes_insertion_idx], unique_classes, node.body[new_classes_insertion_idx:])
+                )
+                node = node.with_changes(body=new_body)
+
         if max_function_index is not None:
             node = node.with_changes(
                 body=(*node.body[: max_function_index + 1], *self.new_functions, *node.body[max_function_index + 1 :])
             )
-        elif class_index is not None:
+        elif max_class_index is not None:
             node = node.with_changes(
-                body=(*node.body[: class_index + 1], *self.new_functions, *node.body[class_index + 1 :])
+                body=(*node.body[: max_class_index + 1], *self.new_functions, *node.body[max_class_index + 1 :])
             )
         else:
             node = node.with_changes(body=(*self.new_functions, *node.body))
@@ -373,18 +399,20 @@ def replace_functions_in_file(
         parsed_function_names.append((class_name, function_name))
 
     # Collect functions we want to modify from the optimized code
-    module = cst.metadata.MetadataWrapper(cst.parse_module(optimized_code))
+    optimized_module = cst.metadata.MetadataWrapper(cst.parse_module(optimized_code))
+    original_module = cst.parse_module(source_code)
+
     visitor = OptimFunctionCollector(preexisting_objects, set(parsed_function_names))
-    module.visit(visitor)
+    optimized_module.visit(visitor)
 
     # Replace these functions in the original code
     transformer = OptimFunctionReplacer(
         modified_functions=visitor.modified_functions,
+        new_classes=visitor.new_classes,
         new_functions=visitor.new_functions,
         new_class_functions=visitor.new_class_functions,
         modified_init_functions=visitor.modified_init_functions,
     )
-    original_module = cst.parse_module(source_code)
     modified_tree = original_module.visit(transformer)
     return modified_tree.code
 

tests/test_code_replacement.py

@@ -215,7 +215,7 @@ def new_function(self, value: cst.Name):
         return other_function(self.name)
     def new_function2(value):
         return value
-    """
+"""
 
     original_code = """import libcst as cst
 from typing import Mandatory
@@ -230,19 +230,28 @@ def other_function(st):
 
 print("Salut monde")
 """
-    expected = """from typing import Mandatory
+    expected = """import libcst as cst
+from typing import Mandatory
+
+class NewClass:
+    def __init__(self, name):
+        self.name = name
+    def new_function(self, value: cst.Name):
+        return other_function(self.name)
+    def new_function2(value):
+        return value
 
 print("Au revoir")
 
 def yet_another_function(values):
     return len(values)
 
-def other_function(st):
-    return(st * 2)
-
 def totally_new_function(value):
     return value
 
+def other_function(st):
+    return(st * 2)
+
 print("Salut monde")
 """
 
@@ -279,7 +288,7 @@ def new_function(self, value):
         return other_function(self.name)
     def new_function2(value):
         return value
-    """
+"""
 
     original_code = """import libcst as cst
 from typing import Mandatory
@@ -296,17 +305,25 @@ def other_function(st):
 """
     expected = """from typing import Mandatory
 
+class NewClass:
+    def __init__(self, name):
+        self.name = name
+    def new_function(self, value):
+        return other_function(self.name)
+    def new_function2(value):
+        return value
+
 print("Au revoir")
 
 def yet_another_function(values):
     return len(values) + 2
 
-def other_function(st):
-    return(st * 2)
-
 def totally_new_function(value):
     return value
 
+def other_function(st):
+    return(st * 2)
+
 print("Salut monde")
 """
 
@@ -3619,4 +3636,110 @@ async def task():
     await asyncio.sleep(1)
     return "done"
 '''
-    assert is_zero_diff(original_code, optimized_code)
+    assert is_zero_diff(original_code, optimized_code)
+
+
+
+def test_code_replacement_with_new_helper_class() -> None:
+    optim_code = """from __future__ import annotations
+
+import itertools
+import re
+from dataclasses import dataclass
+from typing import Any, Callable, Iterator, Sequence
+
+from bokeh.models import HoverTool, Plot, Tool
+
+
+# Move the Item dataclass to module-level to avoid redefining it on every function call
+@dataclass(frozen=True)
+class _RepeatedToolItem:
+    obj: Tool
+    properties: dict[str, Any]
+
+def _collect_repeated_tools(tool_objs: list[Tool]) -> Iterator[Tool]:
+    key: Callable[[Tool], str] = lambda obj: obj.__class__.__name__
+    # Pre-collect properties for all objects by group to avoid repeated calls
+    for _, group in itertools.groupby(sorted(tool_objs, key=key), key=key):
+        grouped = list(group)
+        n = len(grouped)
+        if n > 1:
+            # Precompute all properties once for this group
+            props = [_RepeatedToolItem(obj, obj.properties_with_values()) for obj in grouped]
+            i = 0
+            while i < len(props) - 1:
+                head = props[i]
+                for j in range(i+1, len(props)):
+                    item = props[j]
+                    if item.properties == head.properties:
+                        yield item.obj
+                i += 1
+"""
+
+    original_code = """from __future__ import annotations
+import itertools
+import re
+from bokeh.models import HoverTool, Plot, Tool
+from dataclasses import dataclass
+from typing import Any, Callable, Iterator, Sequence
+
+def _collect_repeated_tools(tool_objs: list[Tool]) -> Iterator[Tool]:
+    @dataclass(frozen=True)
+    class Item:
+        obj: Tool
+        properties: dict[str, Any]
+
+    key: Callable[[Tool], str] = lambda obj: obj.__class__.__name__
+
+    for _, group in itertools.groupby(sorted(tool_objs, key=key), key=key):
+        rest = [ Item(obj, obj.properties_with_values()) for obj in group ]
+        while len(rest) > 1:
+            head, *rest = rest
+            for item in rest:
+                if item.properties == head.properties:
+                    yield item.obj
+"""
+
+    expected = """from __future__ import annotations
+import itertools
+from bokeh.models import Tool
+from dataclasses import dataclass
+from typing import Any, Callable, Iterator
+
+
+# Move the Item dataclass to module-level to avoid redefining it on every function call
+@dataclass(frozen=True)
+class _RepeatedToolItem:
+    obj: Tool
+    properties: dict[str, Any]
+
+def _collect_repeated_tools(tool_objs: list[Tool]) -> Iterator[Tool]:
+    key: Callable[[Tool], str] = lambda obj: obj.__class__.__name__
+    # Pre-collect properties for all objects by group to avoid repeated calls
+    for _, group in itertools.groupby(sorted(tool_objs, key=key), key=key):
+        grouped = list(group)
+        n = len(grouped)
+        if n > 1:
+            # Precompute all properties once for this group
+            props = [_RepeatedToolItem(obj, obj.properties_with_values()) for obj in grouped]
+            i = 0
+            while i < len(props) - 1:
+                head = props[i]
+                for j in range(i+1, len(props)):
+                    item = props[j]
+                    if item.properties == head.properties:
+                        yield item.obj
+                i += 1
+"""
+
+    function_names: list[str] = ["_collect_repeated_tools"]
+    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] = find_preexisting_objects(original_code)
+    new_code: str = replace_functions_and_add_imports(
+        source_code=original_code,
+        function_names=function_names,
+        optimized_code=optim_code,
+        module_abspath=Path(__file__).resolve(),
+        preexisting_objects=preexisting_objects,
+        project_root_path=Path(__file__).resolve().parent.resolve(),
+    )
+    assert new_code == expected