Merge branch 'main' of https://github.com/codeflash-ai/codeflash into part-1-windows-fixes

Author: KRRT7

codeflash/cli_cmds/cmd_init.py

@@ -231,8 +231,7 @@ def collect_setup_info() -> SetupInfo:
     info_panel = Panel(
         Text(
             "📁 Let's identify your Python module directory.\n\n"
-            "This is usually the top-level directory containing all your Python source code.\n"
-            "We've automatically detected some directories for you.",
+            "This is usually the top-level directory containing all your Python source code.\n",
             style="cyan",
         ),
         title="🔍 Module Discovery",
@@ -936,7 +935,6 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
     codeflash_section["module-root"] = setup_info.module_root
     codeflash_section["tests-root"] = setup_info.tests_root
     codeflash_section["test-framework"] = setup_info.test_framework
-    codeflash_section["benchmarks-root"] = setup_info.benchmarks_root if setup_info.benchmarks_root else ""
     codeflash_section["ignore-paths"] = setup_info.ignore_paths
     codeflash_section["disable-telemetry"] = not enable_telemetry
     if setup_info.git_remote not in ["", "origin"]:

codeflash/code_utils/edit_generated_tests.py

@@ -624,45 +624,47 @@ def process_test_files(
                 except Exception as e:
                     logger.debug(str(e))
                     continue
-
-                if not definition or definition[0].type != "function":
-                    continue
-
-                definition_obj = definition[0]
-                definition_path = str(definition_obj.module_path)
-
-                project_root_str = str(project_root_path)
-                if (
-                    definition_path.startswith(project_root_str + os.sep)
-                    and definition_obj.module_name != name.module_name
-                    and definition_obj.full_name is not None
-                ):
-                    # Pre-compute common values outside the inner loop
-                    module_prefix = definition_obj.module_name + "."
-                    full_name_without_module_prefix = definition_obj.full_name.replace(module_prefix, "", 1)
-                    qualified_name_with_modules_from_root = f"{module_name_from_file_path(definition_obj.module_path, project_root_path)}.{full_name_without_module_prefix}"
-
-                    for test_func in test_functions_by_name[scope]:
-                        if test_func.parameters is not None:
-                            if test_framework == "pytest":
-                                scope_test_function = f"{test_func.function_name}[{test_func.parameters}]"
-                            else:  # unittest
-                                scope_test_function = f"{test_func.function_name}_{test_func.parameters}"
-                        else:
-                            scope_test_function = test_func.function_name
-
-                        function_to_test_map[qualified_name_with_modules_from_root].add(
-                            FunctionCalledInTest(
-                                tests_in_file=TestsInFile(
-                                    test_file=test_file,
-                                    test_class=test_func.test_class,
-                                    test_function=scope_test_function,
-                                    test_type=test_func.test_type,
-                                ),
-                                position=CodePosition(line_no=name.line, col_no=name.column),
+                try:
+                    if not definition or definition[0].type != "function":
+                        continue
+                    definition_obj = definition[0]
+                    definition_path = str(definition_obj.module_path)
+
+                    project_root_str = str(project_root_path)
+                    if (
+                        definition_path.startswith(project_root_str + os.sep)
+                        and definition_obj.module_name != name.module_name
+                        and definition_obj.full_name is not None
+                    ):
+                        # Pre-compute common values outside the inner loop
+                        module_prefix = definition_obj.module_name + "."
+                        full_name_without_module_prefix = definition_obj.full_name.replace(module_prefix, "", 1)
+                        qualified_name_with_modules_from_root = f"{module_name_from_file_path(definition_obj.module_path, project_root_path)}.{full_name_without_module_prefix}"
+
+                        for test_func in test_functions_by_name[scope]:
+                            if test_func.parameters is not None:
+                                if test_framework == "pytest":
+                                    scope_test_function = f"{test_func.function_name}[{test_func.parameters}]"
+                                else:  # unittest
+                                    scope_test_function = f"{test_func.function_name}_{test_func.parameters}"
+                            else:
+                                scope_test_function = test_func.function_name
+
+                            function_to_test_map[qualified_name_with_modules_from_root].add(
+                                FunctionCalledInTest(
+                                    tests_in_file=TestsInFile(
+                                        test_file=test_file,
+                                        test_class=test_func.test_class,
+                                        test_function=scope_test_function,
+                                        test_type=test_func.test_type,
+                                    ),
+                                    position=CodePosition(line_no=name.line, col_no=name.column),
+                                )
                             )
-                        )
-                        num_discovered_tests += 1
+                            num_discovered_tests += 1
+                except Exception as e:
+                    logger.debug(str(e))
+                    continue
 
             progress.advance(task_id)
 

codeflash/discovery/discover_unit_tests.py

@@ -1012,15 +1012,20 @@ def find_and_process_best_optimization(
                     optimized_runtime_by_test = (
                         best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case()
                     )
+                    qualified_name = self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)
                     # Add runtime comments to generated tests before creating the PR
                     generated_tests = add_runtime_comments_to_generated_tests(
-                        self.test_cfg, generated_tests, original_runtime_by_test, optimized_runtime_by_test
+                        qualified_name,
+                        self.test_cfg,
+                        generated_tests,
+                        original_runtime_by_test,
+                        optimized_runtime_by_test,
                     )
                     generated_tests_str = "\n\n".join(
                         [test.generated_original_test_source for test in generated_tests.generated_tests]
                     )
                     existing_tests = existing_tests_source_for(
-                        self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root),
+                        qualified_name,
                         function_to_all_tests,
                         test_cfg=self.test_cfg,
                         original_runtimes_all=original_runtime_by_test,

codeflash/optimization/function_optimizer.py

@@ -42,43 +42,39 @@ def existing_tests_source_for(
     rows = []
     headers = ["Test File::Test Function", "Original ⏱️", "Optimized ⏱️", "Speedup"]
     tests_root = test_cfg.tests_root
-    module_root = test_cfg.project_root_path
-    rel_tests_root = tests_root.relative_to(module_root)
     original_tests_to_runtimes: dict[Path, dict[str, int]] = {}
     optimized_tests_to_runtimes: dict[Path, dict[str, int]] = {}
     non_generated_tests = set()
     for test_file in test_files:
-        non_generated_tests.add(Path(test_file.tests_in_file.test_file).relative_to(tests_root))
+        non_generated_tests.add(test_file.tests_in_file.test_file)
     # TODO confirm that original and optimized have the same keys
     all_invocation_ids = original_runtimes_all.keys() | optimized_runtimes_all.keys()
     for invocation_id in all_invocation_ids:
-        rel_path = (
-            Path(invocation_id.test_module_path.replace(".", os.sep)).with_suffix(".py").relative_to(rel_tests_root)
-        )
-        if rel_path not in non_generated_tests:
+        abs_path = Path(invocation_id.test_module_path.replace(".", os.sep)).with_suffix(".py").resolve()
+        if abs_path not in non_generated_tests:
             continue
-        if rel_path not in original_tests_to_runtimes:
-            original_tests_to_runtimes[rel_path] = {}
-        if rel_path not in optimized_tests_to_runtimes:
-            optimized_tests_to_runtimes[rel_path] = {}
+        if abs_path not in original_tests_to_runtimes:
+            original_tests_to_runtimes[abs_path] = {}
+        if abs_path not in optimized_tests_to_runtimes:
+            optimized_tests_to_runtimes[abs_path] = {}
         qualified_name = (
             invocation_id.test_class_name + "." + invocation_id.test_function_name  # type: ignore[operator]
             if invocation_id.test_class_name
             else invocation_id.test_function_name
         )
-        if qualified_name not in original_tests_to_runtimes[rel_path]:
-            original_tests_to_runtimes[rel_path][qualified_name] = 0  # type: ignore[index]
-        if qualified_name not in optimized_tests_to_runtimes[rel_path]:
-            optimized_tests_to_runtimes[rel_path][qualified_name] = 0  # type: ignore[index]
+        if qualified_name not in original_tests_to_runtimes[abs_path]:
+            original_tests_to_runtimes[abs_path][qualified_name] = 0  # type: ignore[index]
+        if qualified_name not in optimized_tests_to_runtimes[abs_path]:
+            optimized_tests_to_runtimes[abs_path][qualified_name] = 0  # type: ignore[index]
         if invocation_id in original_runtimes_all:
-            original_tests_to_runtimes[rel_path][qualified_name] += min(original_runtimes_all[invocation_id])  # type: ignore[index]
+            original_tests_to_runtimes[abs_path][qualified_name] += min(original_runtimes_all[invocation_id])  # type: ignore[index]
         if invocation_id in optimized_runtimes_all:
-            optimized_tests_to_runtimes[rel_path][qualified_name] += min(optimized_runtimes_all[invocation_id])  # type: ignore[index]
+            optimized_tests_to_runtimes[abs_path][qualified_name] += min(optimized_runtimes_all[invocation_id])  # type: ignore[index]
     # parse into string
-    all_rel_paths = (
+    all_abs_paths = (
         original_tests_to_runtimes.keys()
     )  # both will have the same keys as some default values are assigned in the previous loop
-    for filename in sorted(all_rel_paths):
+    for filename in sorted(all_abs_paths):
         all_qualified_names = original_tests_to_runtimes[
             filename
         ].keys()  # both will have the same keys as some default values are assigned in the previous loop
@@ -90,6 +86,7 @@ def existing_tests_source_for(
             ):
                 print_optimized_runtime = format_time(optimized_tests_to_runtimes[filename][qualified_name])
                 print_original_runtime = format_time(original_tests_to_runtimes[filename][qualified_name])
+                print_filename = filename.relative_to(tests_root)
                 greater = (
                     optimized_tests_to_runtimes[filename][qualified_name]
                     > original_tests_to_runtimes[filename][qualified_name]
@@ -104,7 +101,7 @@ def existing_tests_source_for(
                 if greater:
                     rows.append(
                         [
-                            f"`{filename.as_posix()}::{qualified_name}`",
+                            f"`{print_filename.as_posix()}::{qualified_name}`",
                             f"{print_original_runtime}",
                             f"{print_optimized_runtime}",
                             f"⚠️{perf_gain}%",
@@ -113,7 +110,7 @@ def existing_tests_source_for(
                 else:
                     rows.append(
                         [
-                            f"`{filename.as_posix()}::{qualified_name}`",
+                            f"`{print_filename.as_posix()}::{qualified_name}`",
                             f"{print_original_runtime}",
                             f"{print_optimized_runtime}",
                             f"✅{perf_gain}%",