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..1b93016bf --- /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 = "src/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/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 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..2bec158e1 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, @@ -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: diff --git a/tests/test_worktree.py b/tests/test_worktree.py new file mode 100644 index 000000000..b0a66f819 --- /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 / "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