Merge pull request #4 from codeflash-ai/revert-2-function-optimizer-refactor

Revert "Refactor optimizer into FunctionOptimizer class"

Author: alvin-r

codeflash/optimization/function_optimizer.py

@@ -2,6 +2,7 @@
 
 import dataclasses
 import os
+from argparse import Namespace
 from collections import defaultdict
 from pathlib import Path
 
@@ -13,8 +14,7 @@
 )
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import FunctionParent
-from codeflash.optimization.function_optimizer import FunctionOptimizer
-from codeflash.verification.verification_utils import TestConfig
+from codeflash.optimization.optimizer import Optimizer
 
 os.environ["CODEFLASH_API_KEY"] = "cf-test-key"
 
@@ -766,18 +766,24 @@ def main_method(self):
         return HelperClass(self.name).helper_method()
 """
     file_path = Path(__file__).resolve()
+    opt = Optimizer(
+        Namespace(
+            project_root=file_path.parent.resolve(),
+            disable_telemetry=True,
+            tests_root="tests",
+            test_framework="pytest",
+            pytest_cmd="pytest",
+            experiment_id=None,
+            test_project_root=file_path.parent.resolve(),
+        )
+    )
     func_top_optimize = FunctionToOptimize(
         function_name="main_method", file_path=file_path, parents=[FunctionParent("MainClass", "ClassDef")]
     )
-    test_config = TestConfig(
-        tests_root=file_path.parent,
-        tests_project_rootdir=file_path.parent,
-        project_root_path=file_path.parent,
-        test_framework="pytest",
-        pytest_cmd="pytest",
-    )
-    func_optimizer = FunctionOptimizer(function_to_optimize=func_top_optimize, test_cfg=test_config)
-    code_context = func_optimizer.get_code_optimization_context().unwrap()
+    original_code = file_path.read_text()
+    code_context = opt.get_code_optimization_context(
+        function_to_optimize=func_top_optimize, project_root=file_path.parent, original_source_code=original_code
+    ).unwrap()
     assert code_context.code_to_optimize_with_helpers == get_code_output
 
 
@@ -1007,35 +1013,35 @@ def to_name(self) -> str:
 class TestResults(BaseModel):
     def __iter__(self) -> Iterator[FunctionTestInvocation]:
         return iter(self.test_results)
-
+    
     def __len__(self) -> int:
         return len(self.test_results)
-
+    
     def __getitem__(self, index: int) -> FunctionTestInvocation:
         return self.test_results[index]
-
+    
     def __setitem__(self, index: int, value: FunctionTestInvocation) -> None:
         self.test_results[index] = value
-
+    
     def __delitem__(self, index: int) -> None:
         del self.test_results[index]
-
+    
     def __contains__(self, value: FunctionTestInvocation) -> bool:
         return value in self.test_results
-
+    
     def __bool__(self) -> bool:
         return bool(self.test_results)
-
+    
     def __eq__(self, other: object) -> bool:
         # Unordered comparison
         if not isinstance(other, TestResults) or len(self) != len(other):
             return False
-
+        
         # Increase recursion limit only if necessary
         original_recursion_limit = sys.getrecursionlimit()
         if original_recursion_limit < 5000:
             sys.setrecursionlimit(5000)
-
+        
         for test_result in self:
             other_test_result = other.get_by_id(test_result.id)
             if other_test_result is None or not (
@@ -1048,10 +1054,10 @@ def __eq__(self, other: object) -> bool:
             ):
                 sys.setrecursionlimit(original_recursion_limit)
                 return False
-
+        
         sys.setrecursionlimit(original_recursion_limit)
         return True
-
+    
     def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]:
         report = {test_type: {"passed": 0, "failed": 0} for test_type in TestType}
         for test_result in self.test_results:
@@ -1099,8 +1105,8 @@ def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]:
         )
 
     assert (
-            new_code
-            == """from __future__ import annotations
+        new_code
+        == """from __future__ import annotations
 import sys
 from codeflash.verification.comparator import comparator
 from enum import Enum
@@ -1239,21 +1245,21 @@ def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray:
     """Row-wise cosine similarity between two equal-width matrices."""
     if len(X.data) == 0 or len(Y.data) == 0:
         return np.array([])
-
+    
     X_np, Y_np = np.asarray(X.data), np.asarray(Y.data)
     if X_np.shape[1] != Y_np.shape[1]:
         raise ValueError(f"Number of columns in X and Y must be the same. X has shape {X_np.shape} and Y has shape {Y_np.shape}.")
     X_norm = np.linalg.norm(X_np, axis=1, keepdims=True)
     Y_norm = np.linalg.norm(Y_np, axis=1, keepdims=True)
-
+    
     norm_product = X_norm * Y_norm.T
     norm_product[norm_product == 0] = np.inf  # Prevent division by zero
     dot_product = np.dot(X_np, Y_np.T)
     similarity = dot_product / norm_product
-
+    
     # Any NaN or Inf values are set to 0.0
     np.nan_to_num(similarity, copy=False)
-
+    
     return similarity
 def cosine_similarity_top_k(
     X: Matrix,
@@ -1264,15 +1270,15 @@ def cosine_similarity_top_k(
     """Row-wise cosine similarity with optional top-k and score threshold filtering."""
     if len(X.data) == 0 or len(Y.data) == 0:
         return [], []
-
+    
     score_array = cosine_similarity(X, Y)
-
+    
     sorted_idxs = np.argpartition(-score_array.flatten(), range(top_k or len(score_array.flatten())))[:(top_k or len(score_array.flatten()))]
     sorted_idxs = sorted_idxs[score_array.flatten()[sorted_idxs] > (score_threshold if score_threshold is not None else -1)]
-
+    
     ret_idxs = [(x // score_array.shape[1], x % score_array.shape[1]) for x in sorted_idxs]
     scores = score_array.flatten()[sorted_idxs].tolist()
-
+    
     return ret_idxs, scores
 '''
     preexisting_objects: list[tuple[str, list[FunctionParent]]] = find_preexisting_objects(original_code)
@@ -1305,8 +1311,8 @@ def cosine_similarity_top_k(
         project_root_path=Path(__file__).parent.parent.resolve(),
     )
     assert (
-            new_code
-            == '''import numpy as np
+        new_code
+        == '''import numpy as np
 from pydantic.dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 @dataclass(config=dict(arbitrary_types_allowed=True))
@@ -1337,15 +1343,15 @@ def cosine_similarity_top_k(
     """Row-wise cosine similarity with optional top-k and score threshold filtering."""
     if len(X.data) == 0 or len(Y.data) == 0:
         return [], []
-
+    
     score_array = cosine_similarity(X, Y)
-
+    
     sorted_idxs = np.argpartition(-score_array.flatten(), range(top_k or len(score_array.flatten())))[:(top_k or len(score_array.flatten()))]
     sorted_idxs = sorted_idxs[score_array.flatten()[sorted_idxs] > (score_threshold if score_threshold is not None else -1)]
-
+    
     ret_idxs = [(x // score_array.shape[1], x % score_array.shape[1]) for x in sorted_idxs]
     scores = score_array.flatten()[sorted_idxs].tolist()
-
+    
     return ret_idxs, scores
 '''
     )
@@ -1364,8 +1370,8 @@ def cosine_similarity_top_k(
         )
 
     assert (
-            new_helper_code
-            == '''import numpy as np
+        new_helper_code
+        == '''import numpy as np
 from pydantic.dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 @dataclass(config=dict(arbitrary_types_allowed=True))
@@ -1375,21 +1381,21 @@ def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray:
     """Row-wise cosine similarity between two equal-width matrices."""
     if len(X.data) == 0 or len(Y.data) == 0:
         return np.array([])
-
+    
     X_np, Y_np = np.asarray(X.data), np.asarray(Y.data)
     if X_np.shape[1] != Y_np.shape[1]:
         raise ValueError(f"Number of columns in X and Y must be the same. X has shape {X_np.shape} and Y has shape {Y_np.shape}.")
     X_norm = np.linalg.norm(X_np, axis=1, keepdims=True)
     Y_norm = np.linalg.norm(Y_np, axis=1, keepdims=True)
-
+    
     norm_product = X_norm * Y_norm.T
     norm_product[norm_product == 0] = np.inf  # Prevent division by zero
     dot_product = np.dot(X_np, Y_np.T)
     similarity = dot_product / norm_product
-
+    
     # Any NaN or Inf values are set to 0.0
     np.nan_to_num(similarity, copy=False)
-
+    
     return similarity
 def cosine_similarity_top_k(
     X: Matrix,
@@ -1400,15 +1406,15 @@ def cosine_similarity_top_k(
     """Row-wise cosine similarity with optional top-k and score threshold filtering."""
     if len(X.data) == 0 or len(Y.data) == 0:
         return [], []
-
+    
     score_array = cosine_similarity(X, Y)
-
+    
     sorted_idxs = np.argpartition(-score_array.flatten(), range(top_k or len(score_array.flatten())))[:(top_k or len(score_array.flatten()))]
     sorted_idxs = sorted_idxs[score_array.flatten()[sorted_idxs] > (score_threshold if score_threshold is not None else -1)]
-
+    
     ret_idxs = [(x // score_array.shape[1], x % score_array.shape[1]) for x in sorted_idxs]
     scores = score_array.flatten()[sorted_idxs].tolist()
-
+    
     return ret_idxs, scores
 '''
     )
@@ -1475,7 +1481,7 @@ def test_future_aliased_imports_removal() -> None:
 
 def test_0_diff_code_replacement():
     original_code = """from __future__ import annotations
-
+    
 import numpy as np
 def functionA():
     return np.array([1, 2, 3])