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.

Author: KRRT7

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
     )

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:
 # - "<locals>" - 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

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,

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(

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)

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,