From 505691647626e80207f372977a7f4e4912063fae Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 3 Mar 2026 14:47:12 -0800 Subject: [PATCH 01/10] fix: add pre-flight module import check to exit early before API calls When module-root is misconfigured or dependencies are missing, the optimization pipeline previously made expensive API calls (test generation + optimization) before discovering the failure at baseline test execution. This adds a subprocess-based import check in can_be_optimized() that runs before any API calls. If the import fails, it attempts to auto-correct module-root by inferring it from the __init__.py chain and updating pyproject.toml. If auto-correction isn't possible (e.g. missing dependency), it fails early with a clear error message. Co-Authored-By: Claude Opus 4.6 --- codeflash/code_utils/code_utils.py | 51 ++++++++++++++ codeflash/optimization/function_optimizer.py | 72 ++++++++++++++++++++ 2 files changed, 123 insertions(+) diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index 45a64f0fc..9cee209f1 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -7,6 +7,7 @@ import re import shutil import site +import subprocess import sys from contextlib import contextmanager from functools import lru_cache @@ -355,6 +356,56 @@ def module_name_from_file_path(file_path: Path, project_root_path: Path, *, trav raise ValueError(msg) # noqa: B904 +def validate_module_import(module_path: str, project_root: Path) -> tuple[bool, str]: + """Try importing a module in a subprocess to check if it's importable. + + Returns (success, error_message). Uses the same Python executable and + PYTHONPATH setup as the test runner to ensure consistent behavior. + If the import times out (heavy frameworks like TensorFlow), assume it's fine + since the module was found but is just slow to initialize. + """ + from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE + from codeflash.code_utils.shell_utils import make_env_with_project_root + + env = make_env_with_project_root(project_root) + try: + result = subprocess.run( + [SAFE_SYS_EXECUTABLE, "-c", f"import {module_path}"], + env=env, + capture_output=True, + text=True, + timeout=60, + cwd=str(project_root), + errors="replace", + check=False, + ) + except subprocess.TimeoutExpired: + # Heavy modules (e.g. tensorflow) can take a long time — the module exists, just slow to init + return True, "" + if result.returncode == 0: + return True, "" + return False, result.stderr.strip() + + +def infer_module_root_from_file(file_path: Path, pyproject_dir: Path) -> Path | None: + """Infer the correct module-root for a Python file by walking the __init__.py chain. + + Walks up from the file's parent directory toward pyproject_dir. The module-root + is the first ancestor directory that is NOT a Python package (has no __init__.py), + meaning it's the directory that should be on sys.path. + + Returns the inferred module-root path, or None if inference fails. + """ + file_path = file_path.resolve() + pyproject_dir = pyproject_dir.resolve() + current = file_path.parent + while current not in (pyproject_dir, current.parent): + if not (current / "__init__.py").exists(): + return current + current = current.parent + return pyproject_dir + + def file_path_from_module_name(module_name: str, project_root_path: Path) -> Path: """Get file path from module path.""" return project_root_path / (module_name.replace(".", os.sep) + ".py") diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 31b6133d5..2e38c6158 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -42,9 +42,12 @@ extract_unique_errors, file_name_from_test_module_name, get_run_tmp_file, + infer_module_root_from_file, + module_name_from_file_path, normalize_by_max, restore_conftest, unified_diff_strings, + validate_module_import, ) from codeflash.code_utils.config_consts import ( COVERAGE_THRESHOLD, @@ -541,6 +544,61 @@ def parse_line_profile_test_results( # --- End hooks --- + def try_correct_module_root(self) -> bool: + """Try to infer and apply the correct module-root if the current one is wrong. + + Walks the __init__.py chain to determine the correct module-root, validates + it by trying an import, and updates pyproject.toml + in-memory config on success. + """ + from codeflash.cli_cmds.cli import project_root_from_module_root + from codeflash.code_utils.config_parser import find_pyproject_toml + + try: + pyproject_path = find_pyproject_toml(None) + except ValueError: + return False + + pyproject_dir = pyproject_path.parent + inferred_root = infer_module_root_from_file(self.function_to_optimize.file_path, pyproject_dir) + if inferred_root is None or inferred_root.resolve() == self.args.module_root.resolve(): + return False + + new_module_root = inferred_root.resolve() + new_project_root = project_root_from_module_root(new_module_root, pyproject_path) + try: + new_module_path = module_name_from_file_path(self.function_to_optimize.file_path, new_project_root) + except ValueError: + return False + + import_ok, _ = validate_module_import(new_module_path, new_project_root) + if not import_ok: + return False + + # Import succeeded with the inferred module-root — update pyproject.toml + try: + import tomlkit + + with pyproject_path.open("rb") as f: + data = tomlkit.parse(f.read()) + relative_root = os.path.relpath(new_module_root, pyproject_dir) + data["tool"]["codeflash"]["module-root"] = relative_root # type: ignore[index] + with pyproject_path.open("w", encoding="utf-8") as f: + f.write(tomlkit.dumps(data)) + except Exception: + logger.debug("Failed to update pyproject.toml with corrected module-root") + return False + + # Update in-memory config + self.args.module_root = new_module_root + self.args.project_root = new_project_root + self.project_root = new_project_root.resolve() + self.original_module_path = new_module_path + + logger.info( + f"Auto-corrected module-root to '{os.path.relpath(new_module_root, pyproject_dir)}' in pyproject.toml" + ) + return True + def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]: should_run_experiment = self.experiment_id is not None logger.info(f"!lsp|Function Trace ID: {self.function_trace_id}") @@ -555,6 +613,20 @@ def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[P f"Cannot optimize without tests when --no-gen-tests is set." ) + # Pre-flight import check: verify the target module can be imported before expensive API calls + if self.function_to_optimize.language == "python": + import_ok, import_error = validate_module_import(self.original_module_path, self.project_root) + if not import_ok: + corrected = self.try_correct_module_root() + if corrected: + import_ok, import_error = validate_module_import(self.original_module_path, self.project_root) + if not import_ok: + return Failure( + f"Cannot import module '{self.original_module_path}': {import_error}\n" + "This prevents test execution. Please check that all dependencies are installed " + "and that 'module-root' is correctly configured in pyproject.toml." + ) + self.cleanup_leftover_test_return_values() file_name_from_test_module_name.cache_clear() ctx_result = self.get_code_optimization_context() From 916236a0df7e493c6e54f3c9275261126bfd807a Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 3 Mar 2026 14:56:01 -0800 Subject: [PATCH 02/10] fix: correct infer_module_root_from_file to return top-level package The function should return the topmost directory with __init__.py (the top-level package), not the first directory without it. This matches the convention where module-root = package directory and project_root = its parent (the PYTHONPATH entry). Co-Authored-By: Claude Opus 4.6 --- codeflash/code_utils/code_utils.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index 9cee209f1..98bfad2ab 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -390,20 +390,27 @@ def validate_module_import(module_path: str, project_root: Path) -> tuple[bool, def infer_module_root_from_file(file_path: Path, pyproject_dir: Path) -> Path | None: """Infer the correct module-root for a Python file by walking the __init__.py chain. - Walks up from the file's parent directory toward pyproject_dir. The module-root - is the first ancestor directory that is NOT a Python package (has no __init__.py), - meaning it's the directory that should be on sys.path. + Walks up from the file's parent directory toward pyproject_dir, tracking the + topmost directory that contains ``__init__.py`` (i.e. the top-level package). + The module-root is this top-level package directory, since + ``project_root_from_module_root`` will use its parent as the PYTHONPATH entry. Returns the inferred module-root path, or None if inference fails. """ file_path = file_path.resolve() pyproject_dir = pyproject_dir.resolve() current = file_path.parent + top_package: Path | None = None while current not in (pyproject_dir, current.parent): - if not (current / "__init__.py").exists(): - return current + if (current / "__init__.py").exists(): + top_package = current + else: + break current = current.parent - return pyproject_dir + if top_package is not None: + return top_package + # No __init__.py found — treat the file's own directory as the module-root + return file_path.parent def file_path_from_module_name(module_name: str, project_root_path: Path) -> Path: From 96f7200323cb9e46aca9fd4f1d7e4e9ea016be3d Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 3 Mar 2026 15:12:25 -0800 Subject: [PATCH 03/10] fix: always attempt module-root auto-correction before import validation The pre-flight check only tried to correct module-root when the import failed, but imports can succeed even with a misconfigured module-root, causing downstream test result parsing failures after expensive API calls. Co-Authored-By: Claude Opus 4.6 --- codeflash/optimization/function_optimizer.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 2e38c6158..aa4473e32 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -613,19 +613,18 @@ def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[P f"Cannot optimize without tests when --no-gen-tests is set." ) - # Pre-flight import check: verify the target module can be imported before expensive API calls + # Pre-flight: verify module-root consistency and importability before expensive API calls if self.function_to_optimize.language == "python": + # Auto-correct module-root if it doesn't match the inferred root from __init__.py chain + self.try_correct_module_root() + # Now validate the (possibly corrected) module can actually be imported import_ok, import_error = validate_module_import(self.original_module_path, self.project_root) if not import_ok: - corrected = self.try_correct_module_root() - if corrected: - import_ok, import_error = validate_module_import(self.original_module_path, self.project_root) - if not import_ok: - return Failure( - f"Cannot import module '{self.original_module_path}': {import_error}\n" - "This prevents test execution. Please check that all dependencies are installed " - "and that 'module-root' is correctly configured in pyproject.toml." - ) + return Failure( + f"Cannot import module '{self.original_module_path}': {import_error}\n" + "This prevents test execution. Please check that all dependencies are installed " + "and that 'module-root' is correctly configured in pyproject.toml." + ) self.cleanup_leftover_test_return_values() file_name_from_test_module_name.cache_clear() From 4ebf9dfaedf0b0d9e0083768d181583383074d59 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 3 Mar 2026 15:22:54 -0800 Subject: [PATCH 04/10] use find spec instead of subprocess import --- codeflash/code_utils/code_utils.py | 44 +++++++++++++----------------- 1 file changed, 19 insertions(+), 25 deletions(-) diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index 98bfad2ab..d33fb34f9 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -7,7 +7,6 @@ import re import shutil import site -import subprocess import sys from contextlib import contextmanager from functools import lru_cache @@ -357,34 +356,29 @@ def module_name_from_file_path(file_path: Path, project_root_path: Path, *, trav def validate_module_import(module_path: str, project_root: Path) -> tuple[bool, str]: - """Try importing a module in a subprocess to check if it's importable. + """Check if a module is importable using find_spec (no actual import or subprocess). - Returns (success, error_message). Uses the same Python executable and - PYTHONPATH setup as the test runner to ensure consistent behavior. - If the import times out (heavy frameworks like TensorFlow), assume it's fine - since the module was found but is just slow to initialize. + Returns (success, error_message). Uses importlib.util.find_spec to check + module availability without triggering module initialization. """ - from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE - from codeflash.code_utils.shell_utils import make_env_with_project_root + from importlib.util import find_spec - env = make_env_with_project_root(project_root) + project_root_str = str(project_root) + added = False + if project_root_str not in sys.path: + sys.path.insert(0, project_root_str) + added = True try: - result = subprocess.run( - [SAFE_SYS_EXECUTABLE, "-c", f"import {module_path}"], - env=env, - capture_output=True, - text=True, - timeout=60, - cwd=str(project_root), - errors="replace", - check=False, - ) - except subprocess.TimeoutExpired: - # Heavy modules (e.g. tensorflow) can take a long time — the module exists, just slow to init - return True, "" - if result.returncode == 0: - return True, "" - return False, result.stderr.strip() + if find_spec(module_path) is not None: + return True, "" + return False, f"Module '{module_path}' not found (find_spec returned None)" + except ModuleNotFoundError as e: + return False, str(e) + except Exception as e: + return False, f"Error checking module '{module_path}': {e}" + finally: + if added: + sys.path.remove(project_root_str) def infer_module_root_from_file(file_path: Path, pyproject_dir: Path) -> Path | None: From 8d0d8832dd0a55b52bb33d70a0cb56807792feb8 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 3 Mar 2026 15:28:28 -0800 Subject: [PATCH 05/10] test: add unit tests for early module import validation feature Cover validate_module_import, infer_module_root_from_file, and FunctionOptimizer.try_correct_module_root with 21 tests. Co-Authored-By: Claude Opus 4.6 --- .../test_module_import_validation.py | 344 ++++++++++++++++++ 1 file changed, 344 insertions(+) create mode 100644 tests/code_utils/test_module_import_validation.py diff --git a/tests/code_utils/test_module_import_validation.py b/tests/code_utils/test_module_import_validation.py new file mode 100644 index 000000000..658b51000 --- /dev/null +++ b/tests/code_utils/test_module_import_validation.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +import os +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import tomlkit + +from codeflash.code_utils.code_utils import infer_module_root_from_file, validate_module_import + + +class TestValidateModuleImport: + def test_known_stdlib_module(self, tmp_path: Path) -> None: + ok, err = validate_module_import("json", tmp_path) + assert ok is True + assert err == "" + + def test_nonexistent_module(self, tmp_path: Path) -> None: + ok, err = validate_module_import("totally_nonexistent_module_xyz_123", tmp_path) + assert ok is False + assert err != "" + + def test_finds_module_in_project_root(self, tmp_path: Path) -> None: + pkg = tmp_path / "mypkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "utils.py").write_text("x = 1\n", encoding="utf-8") + + ok, err = validate_module_import("mypkg.utils", tmp_path) + assert ok is True + assert err == "" + + def test_project_root_not_left_in_sys_path(self, tmp_path: Path) -> None: + root_str = str(tmp_path) + assert root_str not in sys.path + validate_module_import("nonexistent_mod", tmp_path) + assert root_str not in sys.path + + def test_project_root_preserved_if_already_in_sys_path(self, tmp_path: Path) -> None: + root_str = str(tmp_path) + sys.path.insert(0, root_str) + try: + validate_module_import("json", tmp_path) + assert root_str in sys.path + finally: + sys.path.remove(root_str) + + def test_find_spec_returns_none(self, tmp_path: Path) -> None: + with patch("importlib.util.find_spec", return_value=None) as mock_fs: + ok, err = validate_module_import("some.mod", tmp_path) + mock_fs.assert_called_once_with("some.mod") + assert ok is False + assert "not found" in err + + def test_find_spec_raises_module_not_found(self, tmp_path: Path) -> None: + with patch( + "importlib.util.find_spec", + side_effect=ModuleNotFoundError("No module named 'boom'"), + ): + ok, err = validate_module_import("boom", tmp_path) + assert ok is False + assert "boom" in err + + def test_find_spec_raises_generic_exception(self, tmp_path: Path) -> None: + with patch( + "importlib.util.find_spec", + side_effect=RuntimeError("something broke"), + ): + ok, err = validate_module_import("broken.mod", tmp_path) + assert ok is False + assert "something broke" in err + + def test_sys_path_cleaned_on_exception(self, tmp_path: Path) -> None: + root_str = str(tmp_path) + assert root_str not in sys.path + with patch("importlib.util.find_spec", side_effect=RuntimeError("boom")): + validate_module_import("mod", tmp_path) + assert root_str not in sys.path + + +class TestInferModuleRootFromFile: + def test_single_package(self, tmp_path: Path) -> None: + pkg = tmp_path / "pkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("", encoding="utf-8") + mod = pkg / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + result = infer_module_root_from_file(mod, tmp_path) + assert result is not None + assert result.resolve() == pkg.resolve() + + def test_nested_package_returns_top_level(self, tmp_path: Path) -> None: + pkg = tmp_path / "pkg" + sub = pkg / "sub" + sub.mkdir(parents=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (sub / "__init__.py").write_text("", encoding="utf-8") + mod = sub / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + result = infer_module_root_from_file(mod, tmp_path) + assert result is not None + assert result.resolve() == pkg.resolve() + + def test_deeply_nested_package(self, tmp_path: Path) -> None: + a = tmp_path / "a" + b = a / "b" + c = b / "c" + c.mkdir(parents=True) + for d in (a, b, c): + (d / "__init__.py").write_text("", encoding="utf-8") + mod = c / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + result = infer_module_root_from_file(mod, tmp_path) + assert result is not None + assert result.resolve() == a.resolve() + + def test_no_init_files_returns_parent_dir(self, tmp_path: Path) -> None: + scripts = tmp_path / "scripts" + scripts.mkdir() + mod = scripts / "run.py" + mod.write_text("print('hi')\n", encoding="utf-8") + + result = infer_module_root_from_file(mod, tmp_path) + assert result is not None + assert result.resolve() == scripts.resolve() + + def test_gap_in_init_chain(self, tmp_path: Path) -> None: + outer = tmp_path / "outer" + inner = outer / "inner" + inner.mkdir(parents=True) + (inner / "__init__.py").write_text("", encoding="utf-8") + mod = inner / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + result = infer_module_root_from_file(mod, tmp_path) + assert result is not None + assert result.resolve() == inner.resolve() + + def test_file_directly_in_pyproject_dir(self, tmp_path: Path) -> None: + mod = tmp_path / "standalone.py" + mod.write_text("x = 1\n", encoding="utf-8") + + result = infer_module_root_from_file(mod, tmp_path) + assert result is not None + assert result.resolve() == tmp_path.resolve() + + def test_src_layout(self, tmp_path: Path) -> None: + src = tmp_path / "src" + pkg = src / "pkg" + pkg.mkdir(parents=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + mod = pkg / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + result = infer_module_root_from_file(mod, tmp_path) + assert result is not None + assert result.resolve() == pkg.resolve() + + +class TestTryCorrectModuleRoot: + def _make_optimizer_stub( + self, + file_path: Path, + module_root: Path, + project_root: Path, + original_module_path: str = "pkg.mod", + ) -> MagicMock: + from codeflash.optimization.function_optimizer import FunctionOptimizer + + optimizer = MagicMock(spec=FunctionOptimizer) + optimizer.function_to_optimize = MagicMock() + optimizer.function_to_optimize.file_path = file_path + optimizer.args = MagicMock() + optimizer.args.module_root = module_root + optimizer.args.project_root = project_root + optimizer.project_root = project_root + optimizer.original_module_path = original_module_path + return optimizer + + def test_returns_false_when_pyproject_not_found(self, tmp_path: Path) -> None: + from codeflash.optimization.function_optimizer import FunctionOptimizer + + mod = tmp_path / "pkg" / "mod.py" + mod.parent.mkdir() + mod.write_text("x = 1\n", encoding="utf-8") + optimizer = self._make_optimizer_stub(mod, tmp_path / "pkg", tmp_path) + + with patch( + "codeflash.code_utils.config_parser.find_pyproject_toml", + side_effect=ValueError("not found"), + ): + result = FunctionOptimizer.try_correct_module_root(optimizer) + assert result is False + + def test_returns_false_when_inferred_same_as_current(self, tmp_path: Path) -> None: + from codeflash.optimization.function_optimizer import FunctionOptimizer + + pkg = tmp_path / "pkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("", encoding="utf-8") + mod = pkg / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + pyproject = tmp_path / "pyproject.toml" + pyproject.write_text('[tool.codeflash]\nmodule-root = "pkg"\n', encoding="utf-8") + + optimizer = self._make_optimizer_stub(mod, pkg, tmp_path) + + with patch( + "codeflash.code_utils.config_parser.find_pyproject_toml", + return_value=pyproject, + ): + result = FunctionOptimizer.try_correct_module_root(optimizer) + assert result is False + + def test_corrects_module_root_and_updates_pyproject(self, tmp_path: Path) -> None: + from codeflash.optimization.function_optimizer import FunctionOptimizer + + pkg = tmp_path / "pkg" + sub = pkg / "sub" + sub.mkdir(parents=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (sub / "__init__.py").write_text("", encoding="utf-8") + mod = sub / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + pyproject = tmp_path / "pyproject.toml" + pyproject.write_text('[tool.codeflash]\nmodule-root = "pkg/sub"\n', encoding="utf-8") + + optimizer = self._make_optimizer_stub( + file_path=mod, + module_root=sub, + project_root=tmp_path, + original_module_path="sub.mod", + ) + + with ( + patch( + "codeflash.code_utils.config_parser.find_pyproject_toml", + return_value=pyproject, + ), + patch( + "codeflash.cli_cmds.cli.project_root_from_module_root", + return_value=tmp_path, + ), + patch( + "codeflash.optimization.function_optimizer.module_name_from_file_path", + return_value="pkg.sub.mod", + ), + patch( + "codeflash.optimization.function_optimizer.validate_module_import", + return_value=(True, ""), + ), + ): + result = FunctionOptimizer.try_correct_module_root(optimizer) + + assert result is True + data = tomlkit.parse(pyproject.read_text(encoding="utf-8")) + assert data["tool"]["codeflash"]["module-root"] == os.path.relpath(pkg.resolve(), tmp_path) + + def test_returns_false_when_import_validation_fails(self, tmp_path: Path) -> None: + from codeflash.optimization.function_optimizer import FunctionOptimizer + + pkg = tmp_path / "pkg" + sub = pkg / "sub" + sub.mkdir(parents=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (sub / "__init__.py").write_text("", encoding="utf-8") + mod = sub / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + pyproject = tmp_path / "pyproject.toml" + pyproject.write_text('[tool.codeflash]\nmodule-root = "pkg/sub"\n', encoding="utf-8") + + optimizer = self._make_optimizer_stub( + file_path=mod, + module_root=sub, + project_root=tmp_path, + ) + + with ( + patch( + "codeflash.code_utils.config_parser.find_pyproject_toml", + return_value=pyproject, + ), + patch( + "codeflash.cli_cmds.cli.project_root_from_module_root", + return_value=tmp_path, + ), + patch( + "codeflash.optimization.function_optimizer.module_name_from_file_path", + return_value="pkg.sub.mod", + ), + patch( + "codeflash.optimization.function_optimizer.validate_module_import", + return_value=(False, "Module not found"), + ), + ): + result = FunctionOptimizer.try_correct_module_root(optimizer) + + assert result is False + + def test_returns_false_when_module_name_raises(self, tmp_path: Path) -> None: + from codeflash.optimization.function_optimizer import FunctionOptimizer + + pkg = tmp_path / "pkg" + sub = pkg / "sub" + sub.mkdir(parents=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (sub / "__init__.py").write_text("", encoding="utf-8") + mod = sub / "mod.py" + mod.write_text("x = 1\n", encoding="utf-8") + + pyproject = tmp_path / "pyproject.toml" + pyproject.write_text('[tool.codeflash]\nmodule-root = "pkg/sub"\n', encoding="utf-8") + + optimizer = self._make_optimizer_stub( + file_path=mod, + module_root=sub, + project_root=tmp_path, + ) + + with ( + patch( + "codeflash.code_utils.config_parser.find_pyproject_toml", + return_value=pyproject, + ), + patch( + "codeflash.cli_cmds.cli.project_root_from_module_root", + return_value=tmp_path, + ), + patch( + "codeflash.optimization.function_optimizer.module_name_from_file_path", + side_effect=ValueError("cannot derive module name"), + ), + ): + result = FunctionOptimizer.try_correct_module_root(optimizer) + + assert result is False From 4d21542253ba90ebfe27266862c5c7b6f114f79f Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 3 Mar 2026 15:34:39 -0800 Subject: [PATCH 06/10] only top level imports --- codeflash/code_utils/code_utils.py | 3 +-- codeflash/optimization/function_optimizer.py | 8 +++----- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index d33fb34f9..c9c5c9c1c 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -10,6 +10,7 @@ import sys from contextlib import contextmanager from functools import lru_cache +from importlib.util import find_spec from pathlib import Path from tempfile import TemporaryDirectory @@ -361,8 +362,6 @@ def validate_module_import(module_path: str, project_root: Path) -> tuple[bool, Returns (success, error_message). Uses importlib.util.find_spec to check module availability without triggering module initialization. """ - from importlib.util import find_spec - project_root_str = str(project_root) added = False if project_root_str not in sys.path: diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index aa4473e32..5b5ba1135 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Callable import libcst as cst +import tomlkit from git import Repo as GitRepo from rich.console import Group from rich.panel import Panel @@ -23,6 +24,7 @@ from codeflash.api.aiservice import AiServiceClient, AIServiceRefinerRequest, LocalAiServiceClient from codeflash.api.cfapi import add_code_context_hash, create_staging, get_cfapi_base_urls, mark_optimization_success from codeflash.benchmarking.utils import process_benchmark_data +from codeflash.cli_cmds.cli import project_root_from_module_root from codeflash.cli_cmds.console import ( code_print, console, @@ -61,6 +63,7 @@ EffortLevel, get_effort_value, ) +from codeflash.code_utils.config_parser import find_pyproject_toml from codeflash.code_utils.env_utils import get_pr_number from codeflash.code_utils.formatter import format_code, format_generated_code, sort_imports from codeflash.code_utils.git_utils import git_root_dir @@ -550,9 +553,6 @@ def try_correct_module_root(self) -> bool: Walks the __init__.py chain to determine the correct module-root, validates it by trying an import, and updates pyproject.toml + in-memory config on success. """ - from codeflash.cli_cmds.cli import project_root_from_module_root - from codeflash.code_utils.config_parser import find_pyproject_toml - try: pyproject_path = find_pyproject_toml(None) except ValueError: @@ -576,8 +576,6 @@ def try_correct_module_root(self) -> bool: # Import succeeded with the inferred module-root — update pyproject.toml try: - import tomlkit - with pyproject_path.open("rb") as f: data = tomlkit.parse(f.read()) relative_root = os.path.relpath(new_module_root, pyproject_dir) From 22181a9f5486894f17ba77b3b67f8a42ce722f0f Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 3 Mar 2026 15:47:26 -0800 Subject: [PATCH 07/10] fix: patch mock targets at import site, not definition site Tests were patching functions at their definition modules (e.g. importlib.util.find_spec) instead of where they are imported and used (e.g. codeflash.code_utils.code_utils.find_spec). With `from X import Y`, the local binding is unaffected by patching X.Y. Co-Authored-By: Claude Opus 4.6 --- .../test_module_import_validation.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/code_utils/test_module_import_validation.py b/tests/code_utils/test_module_import_validation.py index 658b51000..b3c67ea2f 100644 --- a/tests/code_utils/test_module_import_validation.py +++ b/tests/code_utils/test_module_import_validation.py @@ -48,7 +48,7 @@ def test_project_root_preserved_if_already_in_sys_path(self, tmp_path: Path) -> sys.path.remove(root_str) def test_find_spec_returns_none(self, tmp_path: Path) -> None: - with patch("importlib.util.find_spec", return_value=None) as mock_fs: + with patch("codeflash.code_utils.code_utils.find_spec", return_value=None) as mock_fs: ok, err = validate_module_import("some.mod", tmp_path) mock_fs.assert_called_once_with("some.mod") assert ok is False @@ -56,7 +56,7 @@ def test_find_spec_returns_none(self, tmp_path: Path) -> None: def test_find_spec_raises_module_not_found(self, tmp_path: Path) -> None: with patch( - "importlib.util.find_spec", + "codeflash.code_utils.code_utils.find_spec", side_effect=ModuleNotFoundError("No module named 'boom'"), ): ok, err = validate_module_import("boom", tmp_path) @@ -65,7 +65,7 @@ def test_find_spec_raises_module_not_found(self, tmp_path: Path) -> None: def test_find_spec_raises_generic_exception(self, tmp_path: Path) -> None: with patch( - "importlib.util.find_spec", + "codeflash.code_utils.code_utils.find_spec", side_effect=RuntimeError("something broke"), ): ok, err = validate_module_import("broken.mod", tmp_path) @@ -75,7 +75,7 @@ def test_find_spec_raises_generic_exception(self, tmp_path: Path) -> None: def test_sys_path_cleaned_on_exception(self, tmp_path: Path) -> None: root_str = str(tmp_path) assert root_str not in sys.path - with patch("importlib.util.find_spec", side_effect=RuntimeError("boom")): + with patch("codeflash.code_utils.code_utils.find_spec", side_effect=RuntimeError("boom")): validate_module_import("mod", tmp_path) assert root_str not in sys.path @@ -191,7 +191,7 @@ def test_returns_false_when_pyproject_not_found(self, tmp_path: Path) -> None: optimizer = self._make_optimizer_stub(mod, tmp_path / "pkg", tmp_path) with patch( - "codeflash.code_utils.config_parser.find_pyproject_toml", + "codeflash.optimization.function_optimizer.find_pyproject_toml", side_effect=ValueError("not found"), ): result = FunctionOptimizer.try_correct_module_root(optimizer) @@ -212,7 +212,7 @@ def test_returns_false_when_inferred_same_as_current(self, tmp_path: Path) -> No optimizer = self._make_optimizer_stub(mod, pkg, tmp_path) with patch( - "codeflash.code_utils.config_parser.find_pyproject_toml", + "codeflash.optimization.function_optimizer.find_pyproject_toml", return_value=pyproject, ): result = FunctionOptimizer.try_correct_module_root(optimizer) @@ -241,11 +241,11 @@ def test_corrects_module_root_and_updates_pyproject(self, tmp_path: Path) -> Non with ( patch( - "codeflash.code_utils.config_parser.find_pyproject_toml", + "codeflash.optimization.function_optimizer.find_pyproject_toml", return_value=pyproject, ), patch( - "codeflash.cli_cmds.cli.project_root_from_module_root", + "codeflash.optimization.function_optimizer.project_root_from_module_root", return_value=tmp_path, ), patch( @@ -285,11 +285,11 @@ def test_returns_false_when_import_validation_fails(self, tmp_path: Path) -> Non with ( patch( - "codeflash.code_utils.config_parser.find_pyproject_toml", + "codeflash.optimization.function_optimizer.find_pyproject_toml", return_value=pyproject, ), patch( - "codeflash.cli_cmds.cli.project_root_from_module_root", + "codeflash.optimization.function_optimizer.project_root_from_module_root", return_value=tmp_path, ), patch( @@ -327,11 +327,11 @@ def test_returns_false_when_module_name_raises(self, tmp_path: Path) -> None: with ( patch( - "codeflash.code_utils.config_parser.find_pyproject_toml", + "codeflash.optimization.function_optimizer.find_pyproject_toml", return_value=pyproject, ), patch( - "codeflash.cli_cmds.cli.project_root_from_module_root", + "codeflash.optimization.function_optimizer.project_root_from_module_root", return_value=tmp_path, ), patch( From 7014a955b1e240eec216baa7ebb3afba45c13194 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Wed, 4 Mar 2026 00:13:47 +0000 Subject: [PATCH 08/10] fix: use as_posix() for paths embedded in Jest reporter JS string template Windows Path backslashes become JS escape sequences when interpolated directly into the script string, mangling the path. Co-authored-by: Aseem Saxena --- tests/test_languages/test_javascript_test_runner.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_languages/test_javascript_test_runner.py b/tests/test_languages/test_javascript_test_runner.py index 9773578fb..d4399516e 100644 --- a/tests/test_languages/test_javascript_test_runner.py +++ b/tests/test_languages/test_javascript_test_runner.py @@ -909,13 +909,13 @@ def test_reporter_produces_valid_junit_xml(self): test_script = Path(tmpdir) / "test_reporter.js" test_script.write_text(f""" // Set env vars BEFORE requiring reporter (matches real Jest behavior) -process.env.JEST_JUNIT_OUTPUT_FILE = '{output_file}'; +process.env.JEST_JUNIT_OUTPUT_FILE = '{output_file.as_posix()}'; process.env.JEST_JUNIT_CLASSNAME = '{{filepath}}'; process.env.JEST_JUNIT_SUITE_NAME = '{{filepath}}'; process.env.JEST_JUNIT_ADD_FILE_ATTRIBUTE = 'true'; process.env.JEST_JUNIT_INCLUDE_CONSOLE_OUTPUT = 'true'; -const Reporter = require('{reporter_path}'); +const Reporter = require('{reporter_path.as_posix()}'); // Mock Jest globalConfig const globalConfig = {{ rootDir: '/tmp/project' }}; From d8986852eafc4cd085f9b86699ddf604a1138ca0 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 3 Mar 2026 18:53:11 -0800 Subject: [PATCH 09/10] refactor: move Python-specific module validation from base class to PythonFunctionOptimizer The try_correct_module_root() method and the Python import validation guard in can_be_optimized() used Python-only concepts (__init__.py chain walking, pyproject.toml updates). Moving them to PythonFunctionOptimizer keeps the base FunctionOptimizer language-agnostic. Co-Authored-By: Claude Opus 4.6 --- .../languages/python/function_optimizer.py | 69 +++++++++++++++++++ codeflash/optimization/function_optimizer.py | 69 ------------------- .../test_module_import_validation.py | 50 +++++++------- 3 files changed, 94 insertions(+), 94 deletions(-) diff --git a/codeflash/languages/python/function_optimizer.py b/codeflash/languages/python/function_optimizer.py index 15babc6b6..01d1b37a3 100644 --- a/codeflash/languages/python/function_optimizer.py +++ b/codeflash/languages/python/function_optimizer.py @@ -1,11 +1,17 @@ from __future__ import annotations import ast +import os from pathlib import Path from typing import TYPE_CHECKING +import tomlkit + +from codeflash.cli_cmds.cli import project_root_from_module_root from codeflash.cli_cmds.console import console, logger +from codeflash.code_utils.code_utils import infer_module_root_from_file, module_name_from_file_path, validate_module_import from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE +from codeflash.code_utils.config_parser import find_pyproject_toml from codeflash.either import Failure, Success from codeflash.languages.python.context.unused_definition_remover import ( detect_unused_helper_functions, @@ -39,6 +45,69 @@ class PythonFunctionOptimizer(FunctionOptimizer): + def try_correct_module_root(self) -> bool: + """Try to infer and apply the correct module-root if the current one is wrong. + + Walks the __init__.py chain to determine the correct module-root, validates + it by trying an import, and updates pyproject.toml + in-memory config on success. + """ + try: + pyproject_path = find_pyproject_toml(None) + except ValueError: + return False + + pyproject_dir = pyproject_path.parent + inferred_root = infer_module_root_from_file(self.function_to_optimize.file_path, pyproject_dir) + if inferred_root is None or inferred_root.resolve() == self.args.module_root.resolve(): + return False + + new_module_root = inferred_root.resolve() + new_project_root = project_root_from_module_root(new_module_root, pyproject_path) + try: + new_module_path = module_name_from_file_path(self.function_to_optimize.file_path, new_project_root) + except ValueError: + return False + + import_ok, _ = validate_module_import(new_module_path, new_project_root) + if not import_ok: + return False + + # Import succeeded with the inferred module-root — update pyproject.toml + try: + with pyproject_path.open("rb") as f: + data = tomlkit.parse(f.read()) + relative_root = os.path.relpath(new_module_root, pyproject_dir) + data["tool"]["codeflash"]["module-root"] = relative_root # type: ignore[index] + with pyproject_path.open("w", encoding="utf-8") as f: + f.write(tomlkit.dumps(data)) + except Exception: + logger.debug("Failed to update pyproject.toml with corrected module-root") + return False + + # Update in-memory config + self.args.module_root = new_module_root + self.args.project_root = new_project_root + self.project_root = new_project_root.resolve() + self.original_module_path = new_module_path + + logger.info( + f"Auto-corrected module-root to '{os.path.relpath(new_module_root, pyproject_dir)}' in pyproject.toml" + ) + return True + + def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]: + # Auto-correct module-root if it doesn't match the inferred root from __init__.py chain + self.try_correct_module_root() + # Validate the (possibly corrected) module can actually be imported + import_ok, import_error = validate_module_import(self.original_module_path, self.project_root) + if not import_ok: + return Failure( + f"Cannot import module '{self.original_module_path}': {import_error}\n" + "This prevents test execution. Please check that all dependencies are installed " + "and that 'module-root' is correctly configured in pyproject.toml." + ) + return super().can_be_optimized() + def get_code_optimization_context(self) -> Result[CodeOptimizationContext, str]: from codeflash.languages.python.context import code_context_extractor diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 5b5ba1135..31b6133d5 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -13,7 +13,6 @@ from typing import TYPE_CHECKING, Callable import libcst as cst -import tomlkit from git import Repo as GitRepo from rich.console import Group from rich.panel import Panel @@ -24,7 +23,6 @@ from codeflash.api.aiservice import AiServiceClient, AIServiceRefinerRequest, LocalAiServiceClient from codeflash.api.cfapi import add_code_context_hash, create_staging, get_cfapi_base_urls, mark_optimization_success from codeflash.benchmarking.utils import process_benchmark_data -from codeflash.cli_cmds.cli import project_root_from_module_root from codeflash.cli_cmds.console import ( code_print, console, @@ -44,12 +42,9 @@ extract_unique_errors, file_name_from_test_module_name, get_run_tmp_file, - infer_module_root_from_file, - module_name_from_file_path, normalize_by_max, restore_conftest, unified_diff_strings, - validate_module_import, ) from codeflash.code_utils.config_consts import ( COVERAGE_THRESHOLD, @@ -63,7 +58,6 @@ EffortLevel, get_effort_value, ) -from codeflash.code_utils.config_parser import find_pyproject_toml from codeflash.code_utils.env_utils import get_pr_number from codeflash.code_utils.formatter import format_code, format_generated_code, sort_imports from codeflash.code_utils.git_utils import git_root_dir @@ -547,56 +541,6 @@ def parse_line_profile_test_results( # --- End hooks --- - def try_correct_module_root(self) -> bool: - """Try to infer and apply the correct module-root if the current one is wrong. - - Walks the __init__.py chain to determine the correct module-root, validates - it by trying an import, and updates pyproject.toml + in-memory config on success. - """ - try: - pyproject_path = find_pyproject_toml(None) - except ValueError: - return False - - pyproject_dir = pyproject_path.parent - inferred_root = infer_module_root_from_file(self.function_to_optimize.file_path, pyproject_dir) - if inferred_root is None or inferred_root.resolve() == self.args.module_root.resolve(): - return False - - new_module_root = inferred_root.resolve() - new_project_root = project_root_from_module_root(new_module_root, pyproject_path) - try: - new_module_path = module_name_from_file_path(self.function_to_optimize.file_path, new_project_root) - except ValueError: - return False - - import_ok, _ = validate_module_import(new_module_path, new_project_root) - if not import_ok: - return False - - # Import succeeded with the inferred module-root — update pyproject.toml - try: - with pyproject_path.open("rb") as f: - data = tomlkit.parse(f.read()) - relative_root = os.path.relpath(new_module_root, pyproject_dir) - data["tool"]["codeflash"]["module-root"] = relative_root # type: ignore[index] - with pyproject_path.open("w", encoding="utf-8") as f: - f.write(tomlkit.dumps(data)) - except Exception: - logger.debug("Failed to update pyproject.toml with corrected module-root") - return False - - # Update in-memory config - self.args.module_root = new_module_root - self.args.project_root = new_project_root - self.project_root = new_project_root.resolve() - self.original_module_path = new_module_path - - logger.info( - f"Auto-corrected module-root to '{os.path.relpath(new_module_root, pyproject_dir)}' in pyproject.toml" - ) - return True - def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]: should_run_experiment = self.experiment_id is not None logger.info(f"!lsp|Function Trace ID: {self.function_trace_id}") @@ -611,19 +555,6 @@ def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[P f"Cannot optimize without tests when --no-gen-tests is set." ) - # Pre-flight: verify module-root consistency and importability before expensive API calls - if self.function_to_optimize.language == "python": - # Auto-correct module-root if it doesn't match the inferred root from __init__.py chain - self.try_correct_module_root() - # Now validate the (possibly corrected) module can actually be imported - import_ok, import_error = validate_module_import(self.original_module_path, self.project_root) - if not import_ok: - return Failure( - f"Cannot import module '{self.original_module_path}': {import_error}\n" - "This prevents test execution. Please check that all dependencies are installed " - "and that 'module-root' is correctly configured in pyproject.toml." - ) - self.cleanup_leftover_test_return_values() file_name_from_test_module_name.cache_clear() ctx_result = self.get_code_optimization_context() diff --git a/tests/code_utils/test_module_import_validation.py b/tests/code_utils/test_module_import_validation.py index b3c67ea2f..66f49015a 100644 --- a/tests/code_utils/test_module_import_validation.py +++ b/tests/code_utils/test_module_import_validation.py @@ -170,9 +170,9 @@ def _make_optimizer_stub( project_root: Path, original_module_path: str = "pkg.mod", ) -> MagicMock: - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer - optimizer = MagicMock(spec=FunctionOptimizer) + optimizer = MagicMock(spec=PythonFunctionOptimizer) optimizer.function_to_optimize = MagicMock() optimizer.function_to_optimize.file_path = file_path optimizer.args = MagicMock() @@ -183,7 +183,7 @@ def _make_optimizer_stub( return optimizer def test_returns_false_when_pyproject_not_found(self, tmp_path: Path) -> None: - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer mod = tmp_path / "pkg" / "mod.py" mod.parent.mkdir() @@ -191,14 +191,14 @@ def test_returns_false_when_pyproject_not_found(self, tmp_path: Path) -> None: optimizer = self._make_optimizer_stub(mod, tmp_path / "pkg", tmp_path) with patch( - "codeflash.optimization.function_optimizer.find_pyproject_toml", + "codeflash.languages.python.function_optimizer.find_pyproject_toml", side_effect=ValueError("not found"), ): - result = FunctionOptimizer.try_correct_module_root(optimizer) + result = PythonFunctionOptimizer.try_correct_module_root(optimizer) assert result is False def test_returns_false_when_inferred_same_as_current(self, tmp_path: Path) -> None: - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer pkg = tmp_path / "pkg" pkg.mkdir() @@ -212,14 +212,14 @@ def test_returns_false_when_inferred_same_as_current(self, tmp_path: Path) -> No optimizer = self._make_optimizer_stub(mod, pkg, tmp_path) with patch( - "codeflash.optimization.function_optimizer.find_pyproject_toml", + "codeflash.languages.python.function_optimizer.find_pyproject_toml", return_value=pyproject, ): - result = FunctionOptimizer.try_correct_module_root(optimizer) + result = PythonFunctionOptimizer.try_correct_module_root(optimizer) assert result is False def test_corrects_module_root_and_updates_pyproject(self, tmp_path: Path) -> None: - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer pkg = tmp_path / "pkg" sub = pkg / "sub" @@ -241,30 +241,30 @@ def test_corrects_module_root_and_updates_pyproject(self, tmp_path: Path) -> Non with ( patch( - "codeflash.optimization.function_optimizer.find_pyproject_toml", + "codeflash.languages.python.function_optimizer.find_pyproject_toml", return_value=pyproject, ), patch( - "codeflash.optimization.function_optimizer.project_root_from_module_root", + "codeflash.languages.python.function_optimizer.project_root_from_module_root", return_value=tmp_path, ), patch( - "codeflash.optimization.function_optimizer.module_name_from_file_path", + "codeflash.languages.python.function_optimizer.module_name_from_file_path", return_value="pkg.sub.mod", ), patch( - "codeflash.optimization.function_optimizer.validate_module_import", + "codeflash.languages.python.function_optimizer.validate_module_import", return_value=(True, ""), ), ): - result = FunctionOptimizer.try_correct_module_root(optimizer) + result = PythonFunctionOptimizer.try_correct_module_root(optimizer) assert result is True data = tomlkit.parse(pyproject.read_text(encoding="utf-8")) assert data["tool"]["codeflash"]["module-root"] == os.path.relpath(pkg.resolve(), tmp_path) def test_returns_false_when_import_validation_fails(self, tmp_path: Path) -> None: - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer pkg = tmp_path / "pkg" sub = pkg / "sub" @@ -285,28 +285,28 @@ def test_returns_false_when_import_validation_fails(self, tmp_path: Path) -> Non with ( patch( - "codeflash.optimization.function_optimizer.find_pyproject_toml", + "codeflash.languages.python.function_optimizer.find_pyproject_toml", return_value=pyproject, ), patch( - "codeflash.optimization.function_optimizer.project_root_from_module_root", + "codeflash.languages.python.function_optimizer.project_root_from_module_root", return_value=tmp_path, ), patch( - "codeflash.optimization.function_optimizer.module_name_from_file_path", + "codeflash.languages.python.function_optimizer.module_name_from_file_path", return_value="pkg.sub.mod", ), patch( - "codeflash.optimization.function_optimizer.validate_module_import", + "codeflash.languages.python.function_optimizer.validate_module_import", return_value=(False, "Module not found"), ), ): - result = FunctionOptimizer.try_correct_module_root(optimizer) + result = PythonFunctionOptimizer.try_correct_module_root(optimizer) assert result is False def test_returns_false_when_module_name_raises(self, tmp_path: Path) -> None: - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer pkg = tmp_path / "pkg" sub = pkg / "sub" @@ -327,18 +327,18 @@ def test_returns_false_when_module_name_raises(self, tmp_path: Path) -> None: with ( patch( - "codeflash.optimization.function_optimizer.find_pyproject_toml", + "codeflash.languages.python.function_optimizer.find_pyproject_toml", return_value=pyproject, ), patch( - "codeflash.optimization.function_optimizer.project_root_from_module_root", + "codeflash.languages.python.function_optimizer.project_root_from_module_root", return_value=tmp_path, ), patch( - "codeflash.optimization.function_optimizer.module_name_from_file_path", + "codeflash.languages.python.function_optimizer.module_name_from_file_path", side_effect=ValueError("cannot derive module name"), ), ): - result = FunctionOptimizer.try_correct_module_root(optimizer) + result = PythonFunctionOptimizer.try_correct_module_root(optimizer) assert result is False From 0404570fcd5704e7bd1a28747e93d2ccdd7bf2d4 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Wed, 4 Mar 2026 02:56:46 +0000 Subject: [PATCH 10/10] fix: guard against None args in try_correct_module_root, fix import ordering Co-authored-by: Aseem Saxena --- codeflash/languages/python/function_optimizer.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/codeflash/languages/python/function_optimizer.py b/codeflash/languages/python/function_optimizer.py index 01d1b37a3..e87880945 100644 --- a/codeflash/languages/python/function_optimizer.py +++ b/codeflash/languages/python/function_optimizer.py @@ -9,7 +9,11 @@ from codeflash.cli_cmds.cli import project_root_from_module_root from codeflash.cli_cmds.console import console, logger -from codeflash.code_utils.code_utils import infer_module_root_from_file, module_name_from_file_path, validate_module_import +from codeflash.code_utils.code_utils import ( + infer_module_root_from_file, + module_name_from_file_path, + validate_module_import, +) from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE from codeflash.code_utils.config_parser import find_pyproject_toml from codeflash.either import Failure, Success @@ -56,6 +60,9 @@ def try_correct_module_root(self) -> bool: except ValueError: return False + if self.args is None: + return False + pyproject_dir = pyproject_path.parent inferred_root = infer_module_root_from_file(self.function_to_optimize.file_path, pyproject_dir) if inferred_root is None or inferred_root.resolve() == self.args.module_root.resolve():