fixes and getting the worktree to work with PR and staging

Author: mohammedahmed18

codeflash/api/cfapi.py

@@ -14,7 +14,7 @@
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
-from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name, git_root_dir
+from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name
 from codeflash.github.PrComment import FileDiffContent, PrComment
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.version import __version__
@@ -206,6 +206,7 @@ def create_staging(
     coverage_message: str,
     replay_tests: str = "",
     concolic_tests: str = "",
+    root_dir: Optional[Path] = None,
 ) -> Response:
     """Create a staging pull request, targeting the specified branch. (usually 'staging').
 
@@ -218,12 +219,10 @@ def create_staging(
     :param coverage_message: Coverage report or summary.
     :return: The response object from the backend.
     """
-    relative_path = explanation.file_path.relative_to(git_root_dir()).as_posix()
+    relative_path = explanation.file_path.relative_to(root_dir).as_posix()
 
     build_file_changes = {
-        Path(p).relative_to(git_root_dir()).as_posix(): FileDiffContent(
-            oldContent=original_code[p], newContent=new_code[p]
-        )
+        Path(p).relative_to(root_dir).as_posix(): FileDiffContent(oldContent=original_code[p], newContent=new_code[p])
         for p in original_code
     }
 

codeflash/code_utils/git_utils.py

@@ -18,7 +18,6 @@
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.compat import codeflash_cache_dir
 from codeflash.code_utils.config_consts import N_CANDIDATES
-from codeflash.lsp.helpers import is_LSP_enabled
 
 if TYPE_CHECKING:
     from git import Repo
@@ -213,19 +212,11 @@ def create_detached_worktree(module_root: Path) -> Optional[Path]:
     current_time_str = time.strftime("%Y%m%d-%H%M%S")
     worktree_dir = worktree_dirs / f"{git_root.name}-{current_time_str}"
 
-    result = subprocess.run(
-        ["git", "worktree", "add", "-d", str(worktree_dir)],
-        cwd=git_root,
-        check=True,
-        stdout=subprocess.DEVNULL if is_LSP_enabled() else None,
-        stderr=subprocess.DEVNULL if is_LSP_enabled() else None,
-    )
-    if result.returncode != 0:
-        logger.error(f"Failed to create worktree: {result.stderr}")
-        return None
+    repository = git.Repo(git_root, search_parent_directories=True)
+
+    repository.git.worktree("add", "-d", str(worktree_dir))
 
     # Get uncommitted diff from the original repo
-    repository = git.Repo(module_root, search_parent_directories=True)
     repository.git.add("-N", ".")  # add the index for untracked files to be included in the diff
     uni_diff_text = repository.git.diff(None, "HEAD", ignore_blank_lines=True, ignore_space_at_eol=True)
 
@@ -234,7 +225,7 @@ def create_detached_worktree(module_root: Path) -> Optional[Path]:
         return worktree_dir
 
     # Write the diff to a temporary file
-    with tempfile.NamedTemporaryFile(mode="w+", suffix=".codeflash.patch", delete=False) as tmp_patch_file:
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".codeflash.patch", delete=False) as tmp_patch_file:
         tmp_patch_file.write(uni_diff_text + "\n")  # the new line here is a must otherwise the last hunk won't be valid
         tmp_patch_file.flush()
 

codeflash/models/models.py

@@ -343,15 +343,12 @@ class TestsInFile:
     test_type: TestType
 
 
-@dataclass(frozen=True)
+@dataclass
 class OptimizedCandidate:
     source_code: CodeStringsMarkdown
     explanation: str
     optimization_id: str
 
-    def set_explanation(self, new_explanation: str) -> None:
-        object.__setattr__(self, "explanation", new_explanation)
-
 
 @dataclass(frozen=True)
 class FunctionCalledInTest:

codeflash/optimization/function_optimizer.py

@@ -1228,7 +1228,7 @@ def process_review(
             benchmark_details=explanation.benchmark_details,
         )
 
-        best_optimization.candidate.set_explanation(new_explanation)
+        best_optimization.candidate.explanation = new_explanation
 
         console.print(Panel(new_explanation_raw_str, title="Best Candidate Explanation", border_style="blue"))
 
@@ -1250,9 +1250,9 @@ def process_review(
 
         if raise_pr and not self.args.staging_review:
             data["git_remote"] = self.args.git_remote
-            check_create_pr(**data)
+            check_create_pr(**data, root_dir=self.project_root)
         elif self.args.staging_review:
-            create_staging(**data)
+            create_staging(**data, root_dir=self.project_root)
         else:
             # Mark optimization success since no PR will be created
             mark_optimization_success(

codeflash/result/create_pr.py

@@ -10,12 +10,7 @@
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_replacer import is_zero_diff
-from codeflash.code_utils.git_utils import (
-    check_and_push_branch,
-    get_current_branch,
-    get_repo_owner_and_name,
-    git_root_dir,
-)
+from codeflash.code_utils.git_utils import check_and_push_branch, get_current_branch, get_repo_owner_and_name
 from codeflash.code_utils.github_utils import github_pr_url
 from codeflash.code_utils.tabulate import tabulate
 from codeflash.code_utils.time_utils import format_perf, format_time
@@ -189,16 +184,17 @@ def check_create_pr(
     replay_tests: str,
     concolic_tests: str,
     git_remote: Optional[str] = None,
+    root_dir: Optional[Path] = None,
 ) -> None:
     pr_number: Optional[int] = env_utils.get_pr_number()
     git_repo = git.Repo(search_parent_directories=True)
 
     if pr_number is not None:
         logger.info(f"Suggesting changes to PR #{pr_number} ...")
         owner, repo = get_repo_owner_and_name(git_repo)
-        relative_path = explanation.file_path.relative_to(git_root_dir()).as_posix()
+        relative_path = explanation.file_path.relative_to(root_dir).as_posix()
         build_file_changes = {
-            Path(p).relative_to(git_root_dir()).as_posix(): FileDiffContent(
+            Path(p).relative_to(root_dir).as_posix(): FileDiffContent(
                 oldContent=original_code[p], newContent=new_code[p]
             )
             for p in original_code
@@ -247,10 +243,10 @@ def check_create_pr(
         if not check_and_push_branch(git_repo, git_remote, wait_for_push=True):
             logger.warning("⏭️ Branch is not pushed, skipping PR creation...")
             return
-        relative_path = explanation.file_path.relative_to(git_root_dir()).as_posix()
+        relative_path = explanation.file_path.relative_to(root_dir).as_posix()
         base_branch = get_current_branch()
         build_file_changes = {
-            Path(p).relative_to(git_root_dir()).as_posix(): FileDiffContent(
+            Path(p).relative_to(root_dir).as_posix(): FileDiffContent(
                 oldContent=original_code[p], newContent=new_code[p]
             )
             for p in original_code