Merge pull request #792 from codeflash-ai/fix/exactly-mimic-args-in-worktree

mirror the args and test_cfg paths for worktree

Author: mohammedahmed18

code_to_optimize/code_directories/nested_module_root/pyproject.toml

@@ -0,0 +1,8 @@
+[tool.codeflash]
+# All paths are relative to this pyproject.toml's directory.
+module-root = "src/app"
+tests-root = "src/tests"
+test-framework = "pytest"
+ignore-paths = []
+disable-telemetry = true
+formatter-cmds = ["disabled"]

code_to_optimize/code_directories/nested_module_root/src/app/main.py

@@ -0,0 +1,10 @@
+def sorter(arr):
+    print("codeflash stdout: Sorting list")
+    for i in range(len(arr)):
+        for j in range(len(arr) - 1):
+            if arr[j] > arr[j + 1]:
+                temp = arr[j]
+                arr[j] = arr[j + 1]
+                arr[j + 1] = temp
+    print(f"result: {arr}")
+    return arr

code_to_optimize/code_directories/nested_module_root/src/tests/.gitkeep

@@ -11,7 +11,6 @@
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.config_parser import parse_config_file
-from codeflash.code_utils.git_utils import git_root_dir
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.version import __version__ as version
 
@@ -223,20 +222,18 @@ def process_pyproject_config(args: Namespace) -> Namespace:
     args.module_root = Path(args.module_root).resolve()
     # If module-root is "." then all imports are relatives to it.
     # in this case, the ".." becomes outside project scope, causing issues with un-importable paths
-    args.project_root = project_root_from_module_root(args.module_root, pyproject_file_path, args.worktree)
+    args.project_root = project_root_from_module_root(args.module_root, pyproject_file_path)
     args.tests_root = Path(args.tests_root).resolve()
     if args.benchmarks_root:
         args.benchmarks_root = Path(args.benchmarks_root).resolve()
-    args.test_project_root = project_root_from_module_root(args.tests_root, pyproject_file_path, args.worktree)
+    args.test_project_root = project_root_from_module_root(args.tests_root, pyproject_file_path)
     if is_LSP_enabled():
         args.all = None
         return args
     return handle_optimize_all_arg_parsing(args)
 
 
-def project_root_from_module_root(module_root: Path, pyproject_file_path: Path, in_worktree: bool = False) -> Path:  # noqa: FBT001, FBT002
-    if in_worktree:
-        return git_root_dir()
+def project_root_from_module_root(module_root: Path, pyproject_file_path: Path) -> Path:
     if pyproject_file_path.parent == module_root:
         return module_root
     return module_root.parent.resolve()

codeflash/cli_cmds/cli.py

@@ -115,24 +115,19 @@ def get_optimizable_functions(
     server: CodeflashLanguageServer, params: OptimizableFunctionsParams
 ) -> dict[str, list[str]]:
     file_path = Path(uris.to_fs_path(params.textDocument.uri))
-    server.show_message_log(f"Getting optimizable functions for: {file_path}", "Info")
     if not server.optimizer:
         return {"status": "error", "message": "optimizer not initialized"}
 
     server.optimizer.args.file = file_path
     server.optimizer.args.function = None  # Always get ALL functions, not just one
     server.optimizer.args.previous_checkpoint_functions = False
 
-    server.show_message_log(f"Calling get_optimizable_functions for {server.optimizer.args.file}...", "Info")
     optimizable_funcs, _, _ = server.optimizer.get_optimizable_functions()
 
     path_to_qualified_names = {}
     for functions in optimizable_funcs.values():
         path_to_qualified_names[file_path] = [func.qualified_name for func in functions]
 
-    server.show_message_log(
-        f"Found {len(path_to_qualified_names)} files with functions: {path_to_qualified_names}", "Info"
-    )
     return path_to_qualified_names
 
 
@@ -177,7 +172,7 @@ def init_project(server: CodeflashLanguageServer, params: ValidateProjectParams)
             else:
                 return {"status": "error", "message": "No pyproject.toml found in workspace."}
 
-    # since we are using worktrees, optimization diffs are generated with respect to the root of the repo, also the args.project_root is set to the root of the repo when creating a worktree
+    # since we are using worktrees, optimization diffs are generated with respect to the root of the repo.
     root = str(git_root_dir())
 
     if getattr(params, "skip_validation", False):

codeflash/lsp/beta.py

@@ -15,7 +15,7 @@
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_utils import cleanup_paths, get_run_tmp_file
 from codeflash.code_utils.env_utils import get_pr_number, is_pr_draft
-from codeflash.code_utils.git_utils import check_running_in_git_repo
+from codeflash.code_utils.git_utils import check_running_in_git_repo, git_root_dir
 from codeflash.code_utils.git_worktree_utils import (
     create_detached_worktree,
     create_diff_patch_from_worktree,
@@ -442,39 +442,50 @@ def worktree_mode(self) -> None:
                 logger.warning("Failed to create worktree. Skipping optimization.")
                 return
             self.current_worktree = worktree_dir
-            self.mutate_args_for_worktree_mode(worktree_dir)
+            self.mirror_paths_for_worktree_mode(worktree_dir)
             # make sure the tests dir is created in the worktree, this can happen if the original tests dir is empty
             Path(self.args.tests_root).mkdir(parents=True, exist_ok=True)
 
-    def mutate_args_for_worktree_mode(self, worktree_dir: Path) -> None:
-        saved_args = copy.deepcopy(self.args)
-        saved_test_cfg = copy.deepcopy(self.test_cfg)
-        self.original_args_and_test_cfg = (saved_args, saved_test_cfg)
-
-        project_root = self.args.project_root
-        module_root = self.args.module_root
-        relative_module_root = module_root.relative_to(project_root)
-        relative_optimized_file = self.args.file.relative_to(project_root) if self.args.file else None
-        relative_tests_root = self.test_cfg.tests_root.relative_to(project_root)
-        relative_benchmarks_root = (
-            self.args.benchmarks_root.relative_to(project_root) if self.args.benchmarks_root else None
+    def mirror_paths_for_worktree_mode(self, worktree_dir: Path) -> None:
+        original_args = copy.deepcopy(self.args)
+        original_test_cfg = copy.deepcopy(self.test_cfg)
+        self.original_args_and_test_cfg = (original_args, original_test_cfg)
+
+        original_git_root = git_root_dir()
+
+        # mirror project_root
+        self.args.project_root = mirror_path(self.args.project_root, original_git_root, worktree_dir)
+        self.test_cfg.project_root_path = mirror_path(self.test_cfg.project_root_path, original_git_root, worktree_dir)
+
+        # mirror module_root
+        self.args.module_root = mirror_path(self.args.module_root, original_git_root, worktree_dir)
+
+        # mirror target file
+        if self.args.file:
+            self.args.file = mirror_path(self.args.file, original_git_root, worktree_dir)
+
+        # mirror tests root
+        self.args.tests_root = mirror_path(self.args.tests_root, original_git_root, worktree_dir)
+        self.test_cfg.tests_root = mirror_path(self.test_cfg.tests_root, original_git_root, worktree_dir)
+
+        # mirror tests project root
+        self.args.test_project_root = mirror_path(self.args.test_project_root, original_git_root, worktree_dir)
+        self.test_cfg.tests_project_rootdir = mirror_path(
+            self.test_cfg.tests_project_rootdir, original_git_root, worktree_dir
         )
 
-        self.args.module_root = worktree_dir / relative_module_root
-        self.args.project_root = worktree_dir
-        self.args.test_project_root = worktree_dir
-        self.args.tests_root = worktree_dir / relative_tests_root
-        if relative_benchmarks_root:
-            self.args.benchmarks_root = worktree_dir / relative_benchmarks_root
-
-        self.test_cfg.project_root_path = worktree_dir
-        self.test_cfg.tests_project_rootdir = worktree_dir
-        self.test_cfg.tests_root = worktree_dir / relative_tests_root
-        if relative_benchmarks_root:
-            self.test_cfg.benchmark_tests_root = worktree_dir / relative_benchmarks_root
-
-        if relative_optimized_file is not None:
-            self.args.file = worktree_dir / relative_optimized_file
+        # mirror benchmarks root paths
+        if self.args.benchmarks_root:
+            self.args.benchmarks_root = mirror_path(self.args.benchmarks_root, original_git_root, worktree_dir)
+        if self.test_cfg.benchmark_tests_root:
+            self.test_cfg.benchmark_tests_root = mirror_path(
+                self.test_cfg.benchmark_tests_root, original_git_root, worktree_dir
+            )
+
+
+def mirror_path(path: Path, src_root: Path, dest_root: Path) -> Path:
+    relative_path = path.relative_to(src_root)
+    return dest_root / relative_path
 
 
 def run_with_args(args: Namespace) -> None:

codeflash/optimization/optimizer.py

@@ -0,0 +1,65 @@
+from argparse import Namespace
+from pathlib import Path
+
+import pytest
+from codeflash.cli_cmds.cli import process_pyproject_config
+from codeflash.optimization.optimizer import Optimizer
+
+
+def test_mirror_paths_for_worktree_mode(monkeypatch: pytest.MonkeyPatch):
+    repo_root = Path(__file__).resolve().parent.parent
+    project_root = repo_root / "code_to_optimize" / "code_directories" / "nested_module_root"
+
+    monkeypatch.setattr("codeflash.optimization.optimizer.git_root_dir", lambda: project_root)
+
+    args = Namespace()
+    args.benchmark = False
+    args.benchmarks_root = None
+
+    args.config_file = project_root / "pyproject.toml"
+    args.file = project_root / "src" / "app" / "main.py"
+    args.worktree = True
+
+    new_args = process_pyproject_config(args)
+
+    optimizer = Optimizer(new_args)
+
+    worktree_dir = repo_root / "worktree"
+    optimizer.mirror_paths_for_worktree_mode(worktree_dir)
+
+    assert optimizer.args.project_root == worktree_dir / "src"
+    assert optimizer.args.test_project_root == worktree_dir / "src"
+    assert optimizer.args.module_root == worktree_dir / "src" / "app"
+    assert optimizer.args.tests_root == worktree_dir / "src" / "tests"
+    assert optimizer.args.file == worktree_dir / "src" / "app" / "main.py"
+
+    assert optimizer.test_cfg.tests_root == worktree_dir / "src" / "tests"
+    assert optimizer.test_cfg.project_root_path == worktree_dir / "src" # same as project_root
+    assert optimizer.test_cfg.tests_project_rootdir == worktree_dir / "src" # same as test_project_root
+
+    # test on our repo
+    monkeypatch.setattr("codeflash.optimization.optimizer.git_root_dir", lambda: repo_root)
+    args = Namespace()
+    args.benchmark = False
+    args.benchmarks_root = None
+
+    args.config_file = repo_root / "pyproject.toml"
+    args.file = repo_root / "codeflash/optimization/optimizer.py"
+    args.worktree = True
+
+    new_args = process_pyproject_config(args)
+
+    optimizer = Optimizer(new_args)
+
+    worktree_dir = repo_root / "worktree"
+    optimizer.mirror_paths_for_worktree_mode(worktree_dir)
+
+    assert optimizer.args.project_root == worktree_dir
+    assert optimizer.args.test_project_root == worktree_dir
+    assert optimizer.args.module_root == worktree_dir / "codeflash"
+    assert optimizer.args.tests_root == worktree_dir / "tests"
+    assert optimizer.args.file == worktree_dir / "codeflash/optimization/optimizer.py"
+
+    assert optimizer.test_cfg.tests_root == worktree_dir / "tests"
+    assert optimizer.test_cfg.project_root_path == worktree_dir # same as project_root
+    assert optimizer.test_cfg.tests_project_rootdir == worktree_dir # same as test_project_root