fix: resolve test file paths in discover_tests_pytest to fix path comparison

Author: KRRT7

codeflash/cli_cmds/cli.py

@@ -352,32 +352,58 @@ def _handle_show_config() -> None:
     from codeflash.setup.detector import detect_project, has_existing_config
 
     project_root = Path.cwd()
-    detected = detect_project(project_root)
-
-    # Check if config exists or is auto-detected
     config_exists, config_file = has_existing_config(project_root)
-    status = "Saved config" if config_exists else "Auto-detected (not saved)"
 
-    console.print()
-    console.print(f"[bold]Codeflash Configuration[/bold] ({status})")
-    if config_exists and config_file:
-        console.print(f"[dim]Config file: {project_root / config_file}[/dim]")
-    console.print()
+    if config_exists:
+        from codeflash.code_utils.config_parser import parse_config_file
 
-    table = Table(show_header=True, header_style="bold cyan")
-    table.add_column("Setting", style="dim")
-    table.add_column("Value")
-
-    table.add_row("Language", detected.language)
-    table.add_row("Project root", str(detected.project_root))
-    table.add_row("Module root", str(detected.module_root))
-    table.add_row("Tests root", str(detected.tests_root) if detected.tests_root else "(not detected)")
-    table.add_row("Test runner", detected.test_runner or "(not detected)")
-    table.add_row("Formatter", ", ".join(detected.formatter_cmds) if detected.formatter_cmds else "(not detected)")
-    table.add_row(
-        "Ignore paths", ", ".join(str(p) for p in detected.ignore_paths) if detected.ignore_paths else "(none)"
-    )
-    table.add_row("Confidence", f"{detected.confidence:.0%}")
+        config, config_file_path = parse_config_file()
+        status = "Saved config"
+
+        console.print()
+        console.print(f"[bold]Codeflash Configuration[/bold] ({status})")
+        console.print(f"[dim]Config file: {config_file_path}[/dim]")
+        console.print()
+
+        table = Table(show_header=True, header_style="bold cyan")
+        table.add_column("Setting", style="dim")
+        table.add_column("Value")
+
+        table.add_row("Project root", str(project_root))
+        table.add_row("Module root", config.get("module_root", "(not set)"))
+        table.add_row("Tests root", config.get("tests_root", "(not set)"))
+        table.add_row("Test runner", config.get("test_framework", config.get("pytest_cmd", "(not set)")))
+        table.add_row(
+            "Formatter",
+            ", ".join(config["formatter_cmds"]) if config.get("formatter_cmds") else "(not set)",
+        )
+        ignore_paths = config.get("ignore_paths", [])
+        table.add_row("Ignore paths", ", ".join(str(p) for p in ignore_paths) if ignore_paths else "(none)")
+    else:
+        detected = detect_project(project_root)
+        status = "Auto-detected (not saved)"
+
+        console.print()
+        console.print(f"[bold]Codeflash Configuration[/bold] ({status})")
+        console.print()
+
+        table = Table(show_header=True, header_style="bold cyan")
+        table.add_column("Setting", style="dim")
+        table.add_column("Value")
+
+        table.add_row("Language", detected.language)
+        table.add_row("Project root", str(detected.project_root))
+        table.add_row("Module root", str(detected.module_root))
+        table.add_row("Tests root", str(detected.tests_root) if detected.tests_root else "(not detected)")
+        table.add_row("Test runner", detected.test_runner or "(not detected)")
+        table.add_row(
+            "Formatter", ", ".join(detected.formatter_cmds) if detected.formatter_cmds else "(not detected)"
+        )
+        table.add_row(
+            "Ignore paths",
+            ", ".join(str(p) for p in detected.ignore_paths) if detected.ignore_paths else "(none)",
+        )
+        table.add_row("Confidence", f"{detected.confidence:.0%}")
 
     console.print(table)
     console.print()

codeflash/code_utils/time_utils.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from codeflash.result.critic import performance_gain
+
 
 def humanize_runtime(time_in_ns: int) -> str:
     runtime_human: str = str(time_in_ns)
@@ -89,3 +91,11 @@ def format_perf(percentage: float) -> str:
     if abs_perc >= 1:
         return f"{percentage:.2f}"
     return f"{percentage:.3f}"
+
+
+def format_runtime_comment(original_time_ns: int, optimized_time_ns: int, comment_prefix: str = "#") -> str:
+    perf_gain = format_perf(
+        abs(performance_gain(original_runtime_ns=original_time_ns, optimized_runtime_ns=optimized_time_ns) * 100)
+    )
+    status = "slower" if optimized_time_ns > original_time_ns else "faster"
+    return f"{comment_prefix} {format_time(original_time_ns)} -> {format_time(optimized_time_ns)} ({perf_gain}% {status})"

codeflash/discovery/discover_unit_tests.py

@@ -728,6 +728,10 @@ def discover_tests_pytest(
         logger.debug(f"Pytest collection exit code: {exitcode}")
     if pytest_rootdir is not None:
         cfg.tests_project_rootdir = Path(pytest_rootdir)
+    if discover_only_these_tests:
+        resolved_discover_only = {p.resolve() for p in discover_only_these_tests}
+    else:
+        resolved_discover_only = None
     file_to_test_map: dict[Path, list[FunctionCalledInTest]] = defaultdict(list)
     for test in tests:
         if "__replay_test" in test["test_file"]:
@@ -737,13 +741,14 @@ def discover_tests_pytest(
         else:
             test_type = TestType.EXISTING_UNIT_TEST
 
+        test_file_path = Path(test["test_file"]).resolve()
         test_obj = TestsInFile(
-            test_file=Path(test["test_file"]),
+            test_file=test_file_path,
             test_class=test["test_class"],
             test_function=test["test_function"],
             test_type=test_type,
         )
-        if discover_only_these_tests and test_obj.test_file not in discover_only_these_tests:
+        if resolved_discover_only and test_obj.test_file not in resolved_discover_only:
             continue
         file_to_test_map[test_obj.test_file].append(test_obj)
     # Within these test files, find the project functions they are referring to and return their names/locations

codeflash/languages/javascript/edit_tests.py

@@ -11,27 +11,8 @@
 from pathlib import Path
 
 from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.time_utils import format_perf, format_time
+from codeflash.code_utils.time_utils import format_runtime_comment
 from codeflash.models.models import GeneratedTests, GeneratedTestsList
-from codeflash.result.critic import performance_gain
-
-
-def format_runtime_comment(original_time: int, optimized_time: int) -> str:
-    """Format a runtime comparison comment for JavaScript.
-
-    Args:
-        original_time: Original runtime in nanoseconds.
-        optimized_time: Optimized runtime in nanoseconds.
-
-    Returns:
-        Formatted comment string with // prefix.
-
-    """
-    perf_gain = format_perf(
-        abs(performance_gain(original_runtime_ns=original_time, optimized_runtime_ns=optimized_time) * 100)
-    )
-    status = "slower" if optimized_time > original_time else "faster"
-    return f"// {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})"
 
 
 def add_runtime_comments(source: str, original_runtimes: dict[str, int], optimized_runtimes: dict[str, int]) -> str:
@@ -120,7 +101,7 @@ def find_matching_test(test_description: str) -> str | None:
             # Only add comment if line has a function call and doesn't already have a comment
             if func_call_pattern.search(line) and "//" not in line and "expect(" in line:
                 orig_time, opt_time = timing_by_full_name[current_matched_full_name]
-                comment = format_runtime_comment(orig_time, opt_time)
+                comment = format_runtime_comment(orig_time, opt_time, comment_prefix="//")
                 logger.debug(f"[js-annotations] Adding comment to test '{current_test_name}': {comment}")
                 # Add comment at end of line
                 line = f"{line.rstrip()}  {comment}"

codeflash/languages/python/context/code_context_extractor.py

@@ -23,10 +23,7 @@
     recurse_sections,
     remove_unused_definitions_by_function_names,
 )
-from codeflash.languages.python.static_analysis.code_extractor import (
-    add_needed_imports_from_module,
-    find_preexisting_objects,
-)
+from codeflash.languages.python.static_analysis.code_extractor import add_needed_imports_from_module, find_preexisting_objects
 from codeflash.models.models import (
     CodeContextType,
     CodeOptimizationContext,
@@ -550,6 +547,35 @@ def get_function_sources_from_jedi(
                             file_path_to_function_source[definition_path].add(function_source)
                             function_source_list.append(function_source)
 
+                    if definition.type == "statement":
+                        try:
+                            for inferred in name.infer():
+                                if (
+                                    inferred.type == "class"
+                                    and inferred.full_name
+                                    and inferred.module_path
+                                    and is_project_path(inferred.module_path, project_root_path)
+                                    and inferred.full_name.startswith(inferred.module_name)
+                                ):
+                                    class_fqn = f"{inferred.full_name}.__init__"
+                                    class_qname = get_qualified_name(inferred.module_name, class_fqn)
+                                    if len(class_qname.split(".")) <= 2:
+                                        class_path = inferred.module_path
+                                        rel = safe_relative_to(class_path, project_root_path)
+                                        if not rel.is_absolute():
+                                            class_path = project_root_path / rel
+                                        class_source = FunctionSource(
+                                            file_path=class_path,
+                                            qualified_name=class_qname,
+                                            fully_qualified_name=class_fqn,
+                                            only_function_name="__init__",
+                                            source_code=inferred.get_line_code(),
+                                        )
+                                        file_path_to_function_source[class_path].add(class_source)
+                                        function_source_list.append(class_source)
+                        except Exception:
+                            logger.debug(f"Error inferring type for statement {definition.full_name}")
+
     return file_path_to_function_source, function_source_list
 
 
@@ -750,6 +776,16 @@ def extract_class_and_bases(
 
             extract_class_and_bases(name, module_path, module_source, module_tree)
 
+            if (module_path, name) not in extracted_classes:
+                for ast_node in module_tree.body:
+                    if isinstance(ast_node, ast.Assign):
+                        for target in ast_node.targets:
+                            if isinstance(target, ast.Name) and target.id == name:
+                                if isinstance(ast_node.value, ast.Call) and isinstance(ast_node.value.func, ast.Name):
+                                    class_name = ast_node.value.func.id
+                                    if class_name not in existing_classes:
+                                        extract_class_and_bases(class_name, module_path, module_source, module_tree)
+
         except Exception:
             logger.debug(f"Error extracting class definition for {name} from {module_name}")
             continue
@@ -759,7 +795,7 @@ def extract_class_and_bases(
         for cls, name in resolve_classes_from_modules(external_base_classes):
             if name in emitted_class_names:
                 continue
-            stub = extract_init_stub(cls, name, require_site_packages=False)
+            stub = extract_class_stub(cls, name, require_site_packages=False)
             if stub is not None:
                 code_strings.append(stub)
                 emitted_class_names.add(name)
@@ -778,7 +814,7 @@ def extract_class_and_bases(
                 continue
             processed_classes.add(cls)
 
-            stub = extract_init_stub(cls, class_name)
+            stub = extract_class_stub(cls, class_name)
             if stub is None:
                 continue
 
@@ -888,22 +924,24 @@ def resolve_transitive_type_deps(cls: type) -> list[type]:
     return deps
 
 
-def extract_init_stub(cls: type, class_name: str, require_site_packages: bool = True) -> CodeString | None:
-    """Extract a stub containing the class definition with only its __init__ method.
+def extract_class_stub(cls: type, class_name: str, require_site_packages: bool = True) -> CodeString | None:
+    """Extract the full class source, falling back to an __init__-only stub.
+
+    Attempts ``inspect.getsource(cls)`` first so the LLM sees every method and
+    attribute.  Falls back to extracting just ``__init__`` when the full source
+    is unavailable (C extensions, dynamically generated classes).  Classes whose
+    ``__init__`` is inherited from ``object`` are still included when the full
+    source can be retrieved.
 
     Args:
-        cls: The class object to extract __init__ from
+        cls: The class object to extract from
         class_name: Name to use for the class in the stub
         require_site_packages: If True, only extract from site-packages. If False, include stdlib too.
 
     """
     import inspect
     import textwrap
 
-    init_method = getattr(cls, "__init__", None)
-    if init_method is None or init_method is object.__init__:
-        return None
-
     try:
         class_file = Path(inspect.getfile(cls))
     except (OSError, TypeError):
@@ -912,17 +950,30 @@ def extract_init_stub(cls: type, class_name: str, require_site_packages: bool =
     if require_site_packages and not path_belongs_to_site_packages(class_file):
         return None
 
+    parts = class_file.parts
+    if "site-packages" in parts:
+        idx = parts.index("site-packages")
+        class_file = Path(*parts[idx + 1 :])
+
+    # Try full class source first
+    try:
+        class_source = inspect.getsource(cls)
+        class_source = textwrap.dedent(class_source)
+        return CodeString(code=class_source, file_path=class_file)
+    except (OSError, TypeError):
+        pass
+
+    # Fallback: __init__-only stub
+    init_method = getattr(cls, "__init__", None)
+    if init_method is None or init_method is object.__init__:
+        return None
+
     try:
         init_source = inspect.getsource(init_method)
         init_source = textwrap.dedent(init_source)
     except (OSError, TypeError):
         return None
 
-    parts = class_file.parts
-    if "site-packages" in parts:
-        idx = parts.index("site-packages")
-        class_file = Path(*parts[idx + 1 :])
-
     class_source = f"class {class_name}:\n" + textwrap.indent(init_source, "    ")
     return CodeString(code=class_source, file_path=class_file)
 
@@ -1080,6 +1131,7 @@ def parse_code_and_prune_cst(
         filtered_node, found_target = prune_cst(
             module,
             target_functions,
+            defs_with_usages=defs_with_usages,
             helpers=helpers_of_helper_functions,
             remove_docstrings=remove_docstrings,
             include_dunder_methods=True,
@@ -1219,7 +1271,7 @@ def prune_cst(
                 stmt,
                 target_functions,
                 class_prefix,
-                defs_with_usages=defs_with_usages,
+                defs_with_usages=None,
                 helpers=helpers,
                 remove_docstrings=remove_docstrings,
                 include_target_in_output=include_target_in_output,