From 10b89fa662703bf605b807dc74d58908f5bd87ae Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 3 Oct 2025 18:19:20 +0300 Subject: [PATCH 1/6] mirror the args in worktree --- codeflash/cli_cmds/cli.py | 9 ++-- codeflash/lsp/beta.py | 7 +--- codeflash/optimization/optimizer.py | 65 ++++++++++++++++++----------- 3 files changed, 45 insertions(+), 36 deletions(-) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 27857806f..933274c19 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -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() diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index ee9e24286..28a6ec655 100644 --- a/codeflash/lsp/beta.py +++ b/codeflash/lsp/beta.py @@ -115,7 +115,6 @@ 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"} @@ -123,16 +122,12 @@ def get_optimizable_functions( 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): diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index c5ab0e1e0..489f075c2 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.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, @@ -447,35 +447,52 @@ def worktree_mode(self) -> None: 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 + 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_module_root = original_args.module_root + original_git_root = git_root_dir().as_posix() + + # mutate project_root + relative_project_root = original_args.project_root.relative_to(original_git_root).as_posix() + # this will be the same as the original project root but in the worktree + new_project_root = worktree_dir / relative_project_root + self.args.project_root = new_project_root + self.test_cfg.project_root_path = new_project_root + + # mutate module_root + relative_module_root = original_module_root.relative_to(original_git_root).as_posix() + self.args.module_root = worktree_dir / relative_module_root + + # mute target file + relative_optimized_file = ( + original_args.file.relative_to(original_git_root).as_posix() if original_args.file else None ) + if relative_optimized_file is not None: + self.args.file = worktree_dir / relative_optimized_file - 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 + # mutate tests root + relative_tests_root = original_test_cfg.tests_root.relative_to(original_git_root).as_posix() + new_tests_root = worktree_dir / relative_tests_root + self.args.tests_root = new_tests_root + self.test_cfg.tests_root = new_tests_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 + # mutate tests project root + relative_tests_project_root = original_args.test_project_root.relative_to(original_git_root).as_posix() + self.args.test_project_root = worktree_dir / relative_tests_project_root + self.test_cfg.tests_project_rootdir = worktree_dir / relative_tests_project_root + + # mutate benchmarks root + relative_benchmarks_root = ( + original_args.benchmarks_root.relative_to(original_git_root).as_posix() + if original_args.benchmarks_root + else None + ) if relative_benchmarks_root: + self.args.benchmarks_root = worktree_dir / 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 - def run_with_args(args: Namespace) -> None: optimizer = None From 85a84773f3488998aef1d3a428d19ba8d4292e5f Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 3 Oct 2025 18:26:47 +0300 Subject: [PATCH 2/6] re-naming --- codeflash/optimization/optimizer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 489f075c2..896715e9f 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -442,11 +442,11 @@ 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: + 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) From f40043209f7dc006ce4870e274e2329078c03fca Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 3 Oct 2025 19:31:07 +0300 Subject: [PATCH 3/6] keep path type --- codeflash/optimization/optimizer.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 896715e9f..f163182b9 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -452,42 +452,38 @@ def mirror_paths_for_worktree_mode(self, worktree_dir: Path) -> None: self.original_args_and_test_cfg = (original_args, original_test_cfg) original_module_root = original_args.module_root - original_git_root = git_root_dir().as_posix() + original_git_root = git_root_dir() # mutate project_root - relative_project_root = original_args.project_root.relative_to(original_git_root).as_posix() + relative_project_root = original_args.project_root.relative_to(original_git_root) # this will be the same as the original project root but in the worktree new_project_root = worktree_dir / relative_project_root self.args.project_root = new_project_root self.test_cfg.project_root_path = new_project_root # mutate module_root - relative_module_root = original_module_root.relative_to(original_git_root).as_posix() + relative_module_root = original_module_root.relative_to(original_git_root) self.args.module_root = worktree_dir / relative_module_root # mute target file - relative_optimized_file = ( - original_args.file.relative_to(original_git_root).as_posix() if original_args.file else None - ) + relative_optimized_file = original_args.file.relative_to(original_git_root) if original_args.file else None if relative_optimized_file is not None: self.args.file = worktree_dir / relative_optimized_file # mutate tests root - relative_tests_root = original_test_cfg.tests_root.relative_to(original_git_root).as_posix() + relative_tests_root = original_test_cfg.tests_root.relative_to(original_git_root) new_tests_root = worktree_dir / relative_tests_root self.args.tests_root = new_tests_root self.test_cfg.tests_root = new_tests_root # mutate tests project root - relative_tests_project_root = original_args.test_project_root.relative_to(original_git_root).as_posix() + relative_tests_project_root = original_args.test_project_root.relative_to(original_git_root) self.args.test_project_root = worktree_dir / relative_tests_project_root self.test_cfg.tests_project_rootdir = worktree_dir / relative_tests_project_root # mutate benchmarks root relative_benchmarks_root = ( - original_args.benchmarks_root.relative_to(original_git_root).as_posix() - if original_args.benchmarks_root - else None + original_args.benchmarks_root.relative_to(original_git_root) if original_args.benchmarks_root else None ) if relative_benchmarks_root: self.args.benchmarks_root = worktree_dir / relative_benchmarks_root From 70d648fb09069bb1463bf73ba5c0120fe251b7a1 Mon Sep 17 00:00:00 2001 From: ali Date: Sat, 4 Oct 2025 23:12:04 +0300 Subject: [PATCH 4/6] tests --- .../nested_module_root/pyproject.toml | 8 +++ .../nested_module_root/src/app/main.py | 10 +++ codeflash/optimization/optimizer.py | 66 +++++++++---------- tests/test_worktree.py | 65 ++++++++++++++++++ 4 files changed, 115 insertions(+), 34 deletions(-) create mode 100644 code_to_optimize/code_directories/nested_module_root/pyproject.toml create mode 100644 code_to_optimize/code_directories/nested_module_root/src/app/main.py create mode 100644 tests/test_worktree.py diff --git a/code_to_optimize/code_directories/nested_module_root/pyproject.toml b/code_to_optimize/code_directories/nested_module_root/pyproject.toml new file mode 100644 index 000000000..34b7b31f3 --- /dev/null +++ b/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 = "tests" +test-framework = "pytest" +ignore-paths = [] +disable-telemetry = true +formatter-cmds = ["disabled"] diff --git a/code_to_optimize/code_directories/nested_module_root/src/app/main.py b/code_to_optimize/code_directories/nested_module_root/src/app/main.py new file mode 100644 index 000000000..9e97f63a0 --- /dev/null +++ b/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 diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index f163182b9..2bec158e1 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -451,43 +451,41 @@ def mirror_paths_for_worktree_mode(self, worktree_dir: Path) -> None: original_test_cfg = copy.deepcopy(self.test_cfg) self.original_args_and_test_cfg = (original_args, original_test_cfg) - original_module_root = original_args.module_root original_git_root = git_root_dir() - # mutate project_root - relative_project_root = original_args.project_root.relative_to(original_git_root) - # this will be the same as the original project root but in the worktree - new_project_root = worktree_dir / relative_project_root - self.args.project_root = new_project_root - self.test_cfg.project_root_path = new_project_root - - # mutate module_root - relative_module_root = original_module_root.relative_to(original_git_root) - self.args.module_root = worktree_dir / relative_module_root - - # mute target file - relative_optimized_file = original_args.file.relative_to(original_git_root) if original_args.file else None - if relative_optimized_file is not None: - self.args.file = worktree_dir / relative_optimized_file - - # mutate tests root - relative_tests_root = original_test_cfg.tests_root.relative_to(original_git_root) - new_tests_root = worktree_dir / relative_tests_root - self.args.tests_root = new_tests_root - self.test_cfg.tests_root = new_tests_root - - # mutate tests project root - relative_tests_project_root = original_args.test_project_root.relative_to(original_git_root) - self.args.test_project_root = worktree_dir / relative_tests_project_root - self.test_cfg.tests_project_rootdir = worktree_dir / relative_tests_project_root - - # mutate benchmarks root - relative_benchmarks_root = ( - original_args.benchmarks_root.relative_to(original_git_root) if original_args.benchmarks_root else None + # 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 ) - if relative_benchmarks_root: - self.args.benchmarks_root = worktree_dir / relative_benchmarks_root - self.test_cfg.benchmark_tests_root = worktree_dir / relative_benchmarks_root + + # 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: diff --git a/tests/test_worktree.py b/tests/test_worktree.py new file mode 100644 index 000000000..22d7ea054 --- /dev/null +++ b/tests/test_worktree.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 + assert optimizer.args.module_root == worktree_dir / "src" / "app" + assert optimizer.args.tests_root == worktree_dir / "tests" + assert optimizer.args.file == worktree_dir / "src" / "app" / "main.py" + + assert optimizer.test_cfg.tests_root == worktree_dir / "tests" + assert optimizer.test_cfg.project_root_path == worktree_dir / "src" # same as project_root + assert optimizer.test_cfg.tests_project_rootdir == worktree_dir # 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 From 7882edce936a98413b0375946a95171e3f4576af Mon Sep 17 00:00:00 2001 From: ali Date: Sat, 4 Oct 2025 23:33:01 +0300 Subject: [PATCH 5/6] also nested tests dir --- .../code_directories/nested_module_root/pyproject.toml | 2 +- tests/test_worktree.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/code_to_optimize/code_directories/nested_module_root/pyproject.toml b/code_to_optimize/code_directories/nested_module_root/pyproject.toml index 34b7b31f3..1b93016bf 100644 --- a/code_to_optimize/code_directories/nested_module_root/pyproject.toml +++ b/code_to_optimize/code_directories/nested_module_root/pyproject.toml @@ -1,7 +1,7 @@ [tool.codeflash] # All paths are relative to this pyproject.toml's directory. module-root = "src/app" -tests-root = "tests" +tests-root = "src/tests" test-framework = "pytest" ignore-paths = [] disable-telemetry = true diff --git a/tests/test_worktree.py b/tests/test_worktree.py index 22d7ea054..b0a66f819 100644 --- a/tests/test_worktree.py +++ b/tests/test_worktree.py @@ -28,14 +28,14 @@ def test_mirror_paths_for_worktree_mode(monkeypatch: pytest.MonkeyPatch): optimizer.mirror_paths_for_worktree_mode(worktree_dir) assert optimizer.args.project_root == worktree_dir / "src" - assert optimizer.args.test_project_root == worktree_dir + 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 / "tests" + 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 / "tests" + 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 # same as test_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) From 6ed7a6c1d3eb9a107be7923bcb196b1c9b360814 Mon Sep 17 00:00:00 2001 From: ali Date: Sat, 4 Oct 2025 23:56:01 +0300 Subject: [PATCH 6/6] keep tests dir --- .../code_directories/nested_module_root/src/tests/.gitkeep | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 code_to_optimize/code_directories/nested_module_root/src/tests/.gitkeep diff --git a/code_to_optimize/code_directories/nested_module_root/src/tests/.gitkeep b/code_to_optimize/code_directories/nested_module_root/src/tests/.gitkeep new file mode 100644 index 000000000..e69de29bb