From 5af9d9c72e248c9d7831d9dd3332fbee6543fcdb Mon Sep 17 00:00:00 2001 From: Codeflash Bot Date: Fri, 17 Oct 2025 21:58:42 +0300 Subject: [PATCH 1/4] correctly resolve test paths for tests runtime comments --- codeflash/code_utils/edit_generated_tests.py | 12 +++++++----- codeflash/lsp/beta.py | 9 +++------ codeflash/optimization/function_optimizer.py | 2 +- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/codeflash/code_utils/edit_generated_tests.py b/codeflash/code_utils/edit_generated_tests.py index abbcb68c1..480ad60d1 100644 --- a/codeflash/code_utils/edit_generated_tests.py +++ b/codeflash/code_utils/edit_generated_tests.py @@ -4,7 +4,7 @@ import os import re from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import libcst as cst from libcst import MetadataWrapper @@ -149,7 +149,7 @@ def leave_SimpleStatementSuite( return updated_node -def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]]) -> dict[str, int]: +def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]], tests_project_rootdir: Path) -> dict[str, int]: unique_inv_ids: dict[str, int] = {} for inv_id, runtimes in inv_id_runtimes.items(): test_qualified_name = ( @@ -157,7 +157,8 @@ def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]]) -> dict[str, i if inv_id.test_class_name else inv_id.test_function_name ) - abs_path = str(Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py").resolve().with_suffix("")) + abs_path = tests_project_rootdir / Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py") + abs_path = str(abs_path.resolve().with_suffix("")) if "__unit_test_" not in abs_path: continue key = test_qualified_name + "#" + abs_path # type: ignore[operator] @@ -174,10 +175,11 @@ def add_runtime_comments_to_generated_tests( generated_tests: GeneratedTestsList, original_runtimes: dict[InvocationId, list[int]], optimized_runtimes: dict[InvocationId, list[int]], + tests_project_rootdir: Optional[Path] = None, ) -> GeneratedTestsList: """Add runtime performance comments to function calls in generated tests.""" - original_runtimes_dict = unique_inv_id(original_runtimes) - optimized_runtimes_dict = unique_inv_id(optimized_runtimes) + original_runtimes_dict = unique_inv_id(original_runtimes, tests_project_rootdir or Path()) + optimized_runtimes_dict = unique_inv_id(optimized_runtimes, tests_project_rootdir or Path()) # Process each generated test modified_tests = [] for test in generated_tests.generated_tests: diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index b93ab197e..9f260f75b 100644 --- a/codeflash/lsp/beta.py +++ b/codeflash/lsp/beta.py @@ -270,15 +270,12 @@ def initialize_function_optimization( if server.optimizer is None: _initialize_optimizer_if_api_key_is_valid(server) - server.optimizer.worktree_mode() - - original_args, _ = server.optimizer.original_args_and_test_cfg - + server.optimizer.args.file = file_path server.optimizer.args.function = params.functionName - original_relative_file_path = file_path.relative_to(original_args.project_root) - server.optimizer.args.file = server.optimizer.current_worktree / original_relative_file_path server.optimizer.args.previous_checkpoint_functions = False + server.optimizer.worktree_mode() + server.show_message_log( f"Args set - function: {server.optimizer.args.function}, file: {server.optimizer.args.file}", "Info" ) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 723a6954d..c30e6040c 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -1375,7 +1375,7 @@ def process_review( ) generated_tests = add_runtime_comments_to_generated_tests( - generated_tests, original_runtime_by_test, optimized_runtime_by_test + generated_tests, original_runtime_by_test, optimized_runtime_by_test, self.test_cfg.tests_project_rootdir ) generated_tests_str = "\n#------------------------------------------------\n".join( From 9a6ef43beabd739c5b17747d84e225359acbab88 Mon Sep 17 00:00:00 2001 From: ali Date: Thu, 23 Oct 2025 19:53:46 +0300 Subject: [PATCH 2/4] fix mypy type error --- codeflash/code_utils/edit_generated_tests.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/codeflash/code_utils/edit_generated_tests.py b/codeflash/code_utils/edit_generated_tests.py index 480ad60d1..1bf42960d 100644 --- a/codeflash/code_utils/edit_generated_tests.py +++ b/codeflash/code_utils/edit_generated_tests.py @@ -158,10 +158,10 @@ def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]], tests_project_ else inv_id.test_function_name ) abs_path = tests_project_rootdir / Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py") - abs_path = str(abs_path.resolve().with_suffix("")) - if "__unit_test_" not in abs_path: + abs_path_str = str(abs_path.resolve()) + if "__unit_test_" not in abs_path_str: continue - key = test_qualified_name + "#" + abs_path # type: ignore[operator] + key = test_qualified_name + "#" + abs_path_str parts = inv_id.iteration_id.split("_").__len__() # type: ignore[union-attr] cur_invid = inv_id.iteration_id.split("_")[0] if parts < 3 else "_".join(inv_id.iteration_id.split("_")[:-1]) # type: ignore[union-attr] match_key = key + "#" + cur_invid From 167a04dd171b5c9703af1ba288ae986987e43c19 Mon Sep 17 00:00:00 2001 From: ali Date: Thu, 23 Oct 2025 19:58:24 +0300 Subject: [PATCH 3/4] fix typing error --- codeflash/code_utils/edit_generated_tests.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/code_utils/edit_generated_tests.py b/codeflash/code_utils/edit_generated_tests.py index 1bf42960d..027415178 100644 --- a/codeflash/code_utils/edit_generated_tests.py +++ b/codeflash/code_utils/edit_generated_tests.py @@ -159,7 +159,7 @@ def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]], tests_project_ ) abs_path = tests_project_rootdir / Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py") abs_path_str = str(abs_path.resolve()) - if "__unit_test_" not in abs_path_str: + if "__unit_test_" not in abs_path_str or not test_qualified_name: continue key = test_qualified_name + "#" + abs_path_str parts = inv_id.iteration_id.split("_").__len__() # type: ignore[union-attr] From 8d160ccc7f232a1240dc91888c1130a03860d161 Mon Sep 17 00:00:00 2001 From: ali Date: Thu, 23 Oct 2025 23:12:13 +0300 Subject: [PATCH 4/4] empty suffix --- codeflash/code_utils/edit_generated_tests.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/code_utils/edit_generated_tests.py b/codeflash/code_utils/edit_generated_tests.py index 027415178..7e8983b3b 100644 --- a/codeflash/code_utils/edit_generated_tests.py +++ b/codeflash/code_utils/edit_generated_tests.py @@ -158,7 +158,7 @@ def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]], tests_project_ else inv_id.test_function_name ) abs_path = tests_project_rootdir / Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py") - abs_path_str = str(abs_path.resolve()) + abs_path_str = str(abs_path.resolve().with_suffix("")) if "__unit_test_" not in abs_path_str or not test_qualified_name: continue key = test_qualified_name + "#" + abs_path_str