split & apply

Author: mohammedahmed18

codeflash/models/models.py

@@ -139,8 +139,11 @@ class CodeString(BaseModel):
     file_path: Optional[Path] = None
 
 
+SPLITTER_MARKER = "# codeflash-splitter__"
+
+
 def get_code_block_splitter(file_path: Path) -> str:
-    return f"# codeflash-splitter__{file_path}"
+    return f"{SPLITTER_MARKER}{file_path}"
 
 
 class CodeStringsMarkdown(BaseModel):
@@ -166,6 +169,20 @@ def markdown(self) -> str:
             ]
         )
 
+    @staticmethod
+    def from_str_with_markers(code_with_markers: str) -> list[CodeString]:
+        pattern = rf"{SPLITTER_MARKER}([^\n]+)\n"
+        matches = list(re.finditer(pattern, code_with_markers))
+
+        results = []
+        for i, match in enumerate(matches):
+            start = match.end()
+            end = matches[i + 1].start() if i + 1 < len(matches) else len(code_with_markers)
+            file_path = match.group(1).strip()
+            code = code_with_markers[start:end].lstrip("\n")
+            results.append(CodeString(file_path=file_path, code=code))
+        return results
+
 
 class CodeOptimizationContext(BaseModel):
     testgen_context_code: str = ""

codeflash/optimization/function_optimizer.py

@@ -62,6 +62,7 @@
 from codeflash.models.models import (
     BestOptimization,
     CodeOptimizationContext,
+    CodeStringsMarkdown,
     GeneratedTests,
     GeneratedTestsList,
     OptimizationSet,
@@ -93,7 +94,6 @@
     from codeflash.either import Result
     from codeflash.models.models import (
         BenchmarkKey,
-        CodeStringsMarkdown,
         CoverageData,
         FunctionCalledInTest,
         FunctionSource,
@@ -621,13 +621,18 @@ def replace_function_and_helpers_with_optimized_code(
         read_writable_functions_by_file_path[self.function_to_optimize.file_path].add(
             self.function_to_optimize.qualified_name
         )
+        code_strings = CodeStringsMarkdown.from_str_with_markers(optimized_code)
+        optimized_code_dict = {code_string.file_path: code_string.code for code_string in code_strings}
+        logger.debug(f"Optimized code: {optimized_code_dict}")
         for helper_function in code_context.helper_functions:
             if helper_function.jedi_definition.type != "class":
                 read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name)
         for module_abspath, qualified_names in read_writable_functions_by_file_path.items():
+            relative_module_path = module_abspath.relative_to(self.project_root)
+            logger.debug(f"applying optimized code to: {relative_module_path}")
             did_update |= replace_function_definitions_in_module(
                 function_names=list(qualified_names),
-                optimized_code=optimized_code,
+                optimized_code=optimized_code_dict.get(relative_module_path),
                 module_abspath=module_abspath,
                 preexisting_objects=code_context.preexisting_objects,
                 project_root_path=self.project_root,