From b6b47ffd8d6e707d3e8f01241fad03e66feadb56 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 22:33:09 -0500 Subject: [PATCH] Unify PYTHONPATH setup into make_env_with_project_root helper Replace duplicated PYTHONPATH-prepending boilerplate across 6 call sites with a single helper in shell_utils. Also adds the missing PYTHONPATH to the concolic test validation subprocess. --- codeflash/benchmarking/trace_benchmarks.py | 9 ++------- codeflash/code_utils/concolic_utils.py | 2 ++ codeflash/code_utils/shell_utils.py | 12 ++++++++++++ codeflash/optimization/function_optimizer.py | 7 ++----- codeflash/tracer.py | 19 +++---------------- codeflash/verification/concolic_testing.py | 10 ++-------- 6 files changed, 23 insertions(+), 36 deletions(-) diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 8217ac37e..98b8e0540 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -1,23 +1,18 @@ from __future__ import annotations -import os import re import subprocess from pathlib import Path from codeflash.cli_cmds.console import logger from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE -from codeflash.code_utils.shell_utils import get_cross_platform_subprocess_run_args +from codeflash.code_utils.shell_utils import get_cross_platform_subprocess_run_args, make_env_with_project_root def trace_benchmarks_pytest( benchmarks_root: Path, tests_root: Path, project_root: Path, trace_file: Path, timeout: int = 300 ) -> None: - benchmark_env = os.environ.copy() - if "PYTHONPATH" not in benchmark_env: - benchmark_env["PYTHONPATH"] = str(project_root) - else: - benchmark_env["PYTHONPATH"] += os.pathsep + str(project_root) + benchmark_env = make_env_with_project_root(project_root) run_args = get_cross_platform_subprocess_run_args( cwd=project_root, env=benchmark_env, timeout=timeout, check=False, text=True, capture_output=True ) diff --git a/codeflash/code_utils/concolic_utils.py b/codeflash/code_utils/concolic_utils.py index 797b4f565..d674be370 100644 --- a/codeflash/code_utils/concolic_utils.py +++ b/codeflash/code_utils/concolic_utils.py @@ -9,6 +9,7 @@ import sentry_sdk from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_temp_dir +from codeflash.code_utils.shell_utils import make_env_with_project_root # Known CrossHair limitations that produce invalid Python syntax in generated tests: # - "" - higher-order functions returning nested functions @@ -37,6 +38,7 @@ def is_valid_concolic_test(test_code: str, project_root: Optional[str] = None) - text=True, cwd=project_root, timeout=10, + env=make_env_with_project_root(project_root) if project_root else None, ) except (subprocess.TimeoutExpired, Exception): return False diff --git a/codeflash/code_utils/shell_utils.py b/codeflash/code_utils/shell_utils.py index df2cff2d6..2052f3e96 100644 --- a/codeflash/code_utils/shell_utils.py +++ b/codeflash/code_utils/shell_utils.py @@ -238,6 +238,18 @@ def save_api_key_to_rc(api_key: str) -> Result[str, str]: ) +def make_env_with_project_root(project_root: Path | str) -> dict[str, str]: + """Return a copy of os.environ with project_root prepended to PYTHONPATH.""" + env = os.environ.copy() + project_root_str = str(project_root) + pythonpath = env.get("PYTHONPATH", "") + if pythonpath: + env["PYTHONPATH"] = f"{project_root_str}{os.pathsep}{pythonpath}" + else: + env["PYTHONPATH"] = project_root_str + return env + + def get_cross_platform_subprocess_run_args( cwd: Path | str | None = None, env: Mapping[str, str] | None = None, diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 6c0283467..519603416 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -69,6 +69,7 @@ from codeflash.code_utils.git_utils import git_root_dir from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test from codeflash.code_utils.line_profile_utils import add_decorator_imports, contains_jit_decorator +from codeflash.code_utils.shell_utils import make_env_with_project_root from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast from codeflash.code_utils.time_utils import humanize_runtime from codeflash.context import code_context_extractor @@ -2821,14 +2822,10 @@ def cleanup_generated_files(self) -> None: def get_test_env( self, codeflash_loop_index: int, codeflash_test_iteration: int, codeflash_tracer_disable: int = 1 ) -> dict: - test_env = os.environ.copy() + test_env = make_env_with_project_root(self.args.project_root) test_env["CODEFLASH_TEST_ITERATION"] = str(codeflash_test_iteration) test_env["CODEFLASH_TRACER_DISABLE"] = str(codeflash_tracer_disable) test_env["CODEFLASH_LOOP_INDEX"] = str(codeflash_loop_index) - if "PYTHONPATH" not in test_env: - test_env["PYTHONPATH"] = str(self.args.project_root) - else: - test_env["PYTHONPATH"] += os.pathsep + str(self.args.project_root) return test_env def line_profiler_step( diff --git a/codeflash/tracer.py b/codeflash/tracer.py index fad0b795d..3f1bde3d1 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -12,7 +12,6 @@ from __future__ import annotations import json -import os import pickle import subprocess import sys @@ -26,6 +25,7 @@ from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE from codeflash.code_utils.config_consts import EffortLevel from codeflash.code_utils.config_parser import parse_config_file +from codeflash.code_utils.shell_utils import make_env_with_project_root from codeflash.tracing.pytest_parallelization import pytest_split if TYPE_CHECKING: @@ -131,13 +131,7 @@ def main(args: Namespace | None = None) -> ArgumentParser: else: updated_sys_argv.append(elem) args_dict["command"] = " ".join(updated_sys_argv) - env = os.environ.copy() - pythonpath = env.get("PYTHONPATH", "") - project_root_str = str(project_root) - if pythonpath: - env["PYTHONPATH"] = f"{project_root_str}{os.pathsep}{pythonpath}" - else: - env["PYTHONPATH"] = project_root_str + env = make_env_with_project_root(project_root) # Disable JIT compilation to ensure tracing captures all function calls env["NUMBA_DISABLE_JIT"] = str(1) env["TORCHDYNAMO_DISABLE"] = str(1) @@ -174,14 +168,7 @@ def main(args: Namespace | None = None) -> ArgumentParser: args_dict["result_pickle_file_path"] = str(result_pickle_file_path) args_dict["command"] = " ".join(sys.argv) - env = os.environ.copy() - # Add project root to PYTHONPATH so imports work correctly - pythonpath = env.get("PYTHONPATH", "") - project_root_str = str(project_root) - if pythonpath: - env["PYTHONPATH"] = f"{project_root_str}{os.pathsep}{pythonpath}" - else: - env["PYTHONPATH"] = project_root_str + env = make_env_with_project_root(project_root) # Disable JIT compilation to ensure tracing captures all function calls env["NUMBA_DISABLE_JIT"] = str(1) env["TORCHDYNAMO_DISABLE"] = str(1) diff --git a/codeflash/verification/concolic_testing.py b/codeflash/verification/concolic_testing.py index 4899b72ff..eda960123 100644 --- a/codeflash/verification/concolic_testing.py +++ b/codeflash/verification/concolic_testing.py @@ -1,7 +1,6 @@ from __future__ import annotations import ast -import os import subprocess import tempfile import time @@ -11,6 +10,7 @@ from codeflash.cli_cmds.console import console, logger from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE from codeflash.code_utils.concolic_utils import clean_concolic_tests, is_valid_concolic_test +from codeflash.code_utils.shell_utils import make_env_with_project_root from codeflash.code_utils.static_analysis import has_typed_parameters from codeflash.discovery.discover_unit_tests import discover_unit_tests from codeflash.languages import is_python @@ -64,13 +64,7 @@ def generate_concolic_tests( logger.info("Generating concolic opcode coverage tests for the original code…") console.rule() try: - env = os.environ.copy() - pythonpath = env.get("PYTHONPATH", "") - project_root_str = str(args.project_root) - if pythonpath: - env["PYTHONPATH"] = f"{project_root_str}{os.pathsep}{pythonpath}" - else: - env["PYTHONPATH"] = project_root_str + env = make_env_with_project_root(args.project_root) cover_result = subprocess.run( [ SAFE_SYS_EXECUTABLE,