Merge remote-tracking branch 'origin/main' into cf-773

Author: Codeflash Bot

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
 

codeflash/code_utils/formatter.py

@@ -8,7 +8,7 @@
 import subprocess
 import tempfile
 from pathlib import Path
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 import isort
 
@@ -178,10 +178,10 @@ def format_code(
         return formatted_code
 
 
-def sort_imports(code: str, *, float_to_top: bool = False) -> str:
+def sort_imports(code: str, **kwargs: Any) -> str:  # noqa : ANN401
     try:
         # Deduplicate and sort imports, modify the code in memory, not on disk
-        sorted_code = isort.code(code=code, float_to_top=float_to_top)
+        sorted_code = isort.code(code, **kwargs)
     except Exception:  # this will also catch the FileSkipComment exception, use this fn everywhere
         logger.exception("Failed to sort imports with isort.")
         return code  # Fall back to original code if isort fails

codeflash/code_utils/instrument_existing_tests.py

@@ -6,7 +6,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-import isort
 import libcst as cst
 
 from codeflash.cli_cmds.console import logger
@@ -741,7 +740,7 @@ def inject_async_profiling_into_existing_test(
         new_imports.append(ast.Import(names=[ast.alias(name="timeout_decorator")]))
 
     tree.body = [*new_imports, *tree.body]
-    return True, isort.code(ast.unparse(tree), float_to_top=True)
+    return True, sort_imports(ast.unparse(tree), float_to_top=True)
 
 
 def inject_profiling_into_existing_test(
@@ -789,7 +788,7 @@ def inject_profiling_into_existing_test(
     additional_functions = [create_wrapper_function(mode)]
 
     tree.body = [*new_imports, *additional_functions, *tree.body]
-    return True, isort.code(ast.unparse(tree), float_to_top=True)
+    return True, sort_imports(ast.unparse(tree), float_to_top=True)
 
 
 def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR) -> ast.FunctionDef: