refactor compat.py

Author: KRRT7

codeflash/benchmarking/trace_benchmarks.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 
 from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
+from codeflash.code_utils.compat import get_safe_sys_executable
 from codeflash.code_utils.shell_utils import get_cross_platform_subprocess_run_args
 
 
@@ -23,7 +23,7 @@ def trace_benchmarks_pytest(
     )
     result = subprocess.run(  # noqa: PLW1510
         [
-            SAFE_SYS_EXECUTABLE,
+            get_safe_sys_executable(),
             Path(__file__).parent / "pytest_new_process_trace_benchmarks.py",
             benchmarks_root,
             tests_root,

codeflash/code_utils/checkpoint.py

@@ -10,7 +10,7 @@
 from rich.prompt import Confirm
 
 from codeflash.cli_cmds.console import console
-from codeflash.code_utils.compat import codeflash_temp_dir
+from codeflash.code_utils.compat import get_codeflash_temp_dir
 
 if TYPE_CHECKING:
     import argparse
@@ -19,7 +19,7 @@
 class CodeflashRunCheckpoint:
     def __init__(self, module_root: Path, checkpoint_dir: Path | None = None) -> None:
         if checkpoint_dir is None:
-            checkpoint_dir = codeflash_temp_dir
+            checkpoint_dir = get_codeflash_temp_dir()
         self.module_root = module_root
         self.checkpoint_dir = Path(checkpoint_dir)
         # Create a unique checkpoint file name
@@ -141,8 +141,8 @@ def get_all_historical_functions(module_root: Path, checkpoint_dir: Path) -> dic
 
 def ask_should_use_checkpoint_get_functions(args: argparse.Namespace) -> Optional[dict[str, dict[str, str]]]:
     previous_checkpoint_functions = None
-    if args.all and codeflash_temp_dir.is_dir():
-        previous_checkpoint_functions = get_all_historical_functions(args.module_root, codeflash_temp_dir)
+    if args.all and get_codeflash_temp_dir().is_dir():
+        previous_checkpoint_functions = get_all_historical_functions(args.module_root, get_codeflash_temp_dir())
         if previous_checkpoint_functions and Confirm.ask(
             "Previous Checkpoint detected from an incomplete optimization run, shall I continue the optimization from that point?",
             default=True,

codeflash/code_utils/compat.py

@@ -13,10 +13,6 @@
 if TYPE_CHECKING:
     from jedi.api.environment import InterpreterEnvironment
 
-    codeflash_temp_dir: Path
-    codeflash_cache_dir: Path
-    codeflash_cache_db: Path
-
 
 def is_compiled_or_bundled_binary() -> bool:
     """Check if running in a compiled/bundled binary."""
@@ -85,41 +81,34 @@ def _find_python_executable() -> str:
     return sys.executable
 
 
-class Compat:
-    # os-independent newline
-    LF: str = os.linesep
-
-    IS_POSIX: bool = os.name != "nt"
+LF: str = os.linesep
+IS_POSIX: bool = os.name != "nt"
 
-    @property
-    def SAFE_SYS_EXECUTABLE(self) -> str:
-        return Path(_find_python_executable()).as_posix()
 
-    @property
-    def codeflash_cache_dir(self) -> Path:
-        return Path(user_config_dir(appname="codeflash", appauthor="codeflash-ai", ensure_exists=True))
+@lru_cache(maxsize=1)
+def get_safe_sys_executable() -> str:
+    """Get a safe Python executable path with forward slashes."""
+    return Path(_find_python_executable()).as_posix()
 
-    @property
-    def codeflash_temp_dir(self) -> Path:
-        temp_dir = Path(tempfile.gettempdir()) / "codeflash"
-        if not temp_dir.exists():
-            temp_dir.mkdir(parents=True, exist_ok=True)
-        return temp_dir
 
-    @property
-    def codeflash_cache_db(self) -> Path:
-        return self.codeflash_cache_dir / "codeflash_cache.db"
+@lru_cache(maxsize=1)
+def get_codeflash_cache_dir() -> Path:
+    """Get the codeflash cache directory, creating it if necessary."""
+    return Path(user_config_dir(appname="codeflash", appauthor="codeflash-ai", ensure_exists=True))
 
 
-_compat = Compat()
+@lru_cache(maxsize=1)
+def get_codeflash_temp_dir() -> Path:
+    """Get the codeflash temp directory, creating it if necessary."""
+    temp_dir = Path(tempfile.gettempdir()) / "codeflash"
+    temp_dir.mkdir(parents=True, exist_ok=True)
+    return temp_dir
 
 
-codeflash_temp_dir = _compat.codeflash_temp_dir
-codeflash_cache_dir = _compat.codeflash_cache_dir
-codeflash_cache_db = _compat.codeflash_cache_db
-LF = _compat.LF
-SAFE_SYS_EXECUTABLE = _compat.SAFE_SYS_EXECUTABLE
-IS_POSIX = _compat.IS_POSIX
+@lru_cache(maxsize=1)
+def get_codeflash_cache_db() -> Path:
+    """Get the path to the codeflash cache database."""
+    return get_codeflash_cache_dir() / "codeflash_cache.db"
 
 
 @lru_cache(maxsize=1)

codeflash/code_utils/concolic_utils.py

@@ -8,7 +8,7 @@
 
 import sentry_sdk
 
-from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_temp_dir
+from codeflash.code_utils.compat import get_codeflash_temp_dir, get_safe_sys_executable
 
 
 def is_valid_concolic_test(test_code: str, project_root: Optional[str] = None) -> bool:
@@ -18,12 +18,12 @@ def is_valid_concolic_test(test_code: str, project_root: Optional[str] = None) -
         sentry_sdk.capture_message(f"CrossHair generated test with syntax error:\n{test_code}")
         return False
 
-    temp_path = (codeflash_temp_dir / f"concolic_test_{uuid.uuid4().hex}.py").resolve()
+    temp_path = (get_codeflash_temp_dir() / f"concolic_test_{uuid.uuid4().hex}.py").resolve()
     temp_path.write_text(test_code, encoding="utf-8")
 
     try:
         result = subprocess.run(
-            [SAFE_SYS_EXECUTABLE, "-m", "pytest", "--collect-only", "-q", temp_path.as_posix()],
+            [get_safe_sys_executable(), "-m", "pytest", "--collect-only", "-q", temp_path.as_posix()],
             check=False,
             capture_output=True,
             text=True,

codeflash/code_utils/git_worktree_utils.py

@@ -12,11 +12,11 @@
 import git
 
 from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.compat import codeflash_cache_dir
+from codeflash.code_utils.compat import get_codeflash_cache_dir
 from codeflash.code_utils.git_utils import check_running_in_git_repo, git_root_dir
 
-worktree_dirs = codeflash_cache_dir / "worktrees"
-patches_dir = codeflash_cache_dir / "patches"
+worktree_dirs = get_codeflash_cache_dir() / "worktrees"
+patches_dir = get_codeflash_cache_dir() / "patches"
 
 
 def create_worktree_snapshot_commit(worktree_dir: Path, commit_message: str) -> None:

codeflash/discovery/discover_unit_tests.py

@@ -28,7 +28,7 @@
     get_run_tmp_file,
     module_name_from_file_path,
 )
-from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_cache_db, is_compiled_or_bundled_binary
+from codeflash.code_utils.compat import get_codeflash_cache_db, get_safe_sys_executable, is_compiled_or_bundled_binary
 from codeflash.code_utils.shell_utils import get_cross_platform_subprocess_run_args
 from codeflash.models.models import CodePosition, FunctionCalledInTest, TestsInFile, TestType
 
@@ -137,7 +137,7 @@ class TestsCache:
 
     def __init__(self, project_root_path: str | Path) -> None:
         self.project_root_path = Path(project_root_path).resolve().as_posix()
-        self.connection = sqlite3.connect(codeflash_cache_db)
+        self.connection = sqlite3.connect(get_codeflash_cache_db())
         self.cur = self.connection.cursor()
 
         self.cur.execute(
@@ -675,7 +675,13 @@ def discover_tests_pytest(
             cwd=project_root, check=False, text=True, capture_output=True
         )
         result = subprocess.run(  # noqa: PLW1510
-            [SAFE_SYS_EXECUTABLE, str(discovery_script_path), str(project_root), str(tests_root), str(tmp_pickle_path)],
+            [
+                get_safe_sys_executable(),
+                str(discovery_script_path),
+                str(project_root),
+                str(tests_root),
+                str(tmp_pickle_path),
+            ],
             **run_kwargs,
         )
     try:

codeflash/tracer.py

@@ -23,7 +23,7 @@
 from codeflash.cli_cmds.cli import project_root_from_module_root
 from codeflash.cli_cmds.console import console
 from codeflash.code_utils.code_utils import get_run_tmp_file
-from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
+from codeflash.code_utils.compat import get_safe_sys_executable
 from codeflash.code_utils.config_consts import EffortLevel
 from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.tracing.pytest_parallelization import pytest_split
@@ -141,7 +141,7 @@ def main(args: Namespace | None = None) -> ArgumentParser:
                     processes.append(
                         subprocess.Popen(
                             [
-                                SAFE_SYS_EXECUTABLE,
+                                get_safe_sys_executable(),
                                 Path(__file__).parent / "tracing" / "tracing_new_process.py",
                                 *updated_sys_argv,
                                 json.dumps(args_dict),
@@ -178,7 +178,7 @@ def main(args: Namespace | None = None) -> ArgumentParser:
 
                 subprocess.run(
                     [
-                        SAFE_SYS_EXECUTABLE,
+                        get_safe_sys_executable(),
                         Path(__file__).parent / "tracing" / "tracing_new_process.py",
                         *sys.argv,
                         json.dumps(args_dict),

codeflash/verification/concolic_testing.py

@@ -8,7 +8,7 @@
 from typing import TYPE_CHECKING
 
 from codeflash.cli_cmds.console import console, logger
-from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
+from codeflash.code_utils.compat import get_safe_sys_executable
 from codeflash.code_utils.concolic_utils import clean_concolic_tests, is_valid_concolic_test
 from codeflash.code_utils.static_analysis import has_typed_parameters
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
@@ -44,7 +44,7 @@ def generate_concolic_tests(
         try:
             cover_result = subprocess.run(
                 [
-                    SAFE_SYS_EXECUTABLE,
+                    get_safe_sys_executable(),
                     "-m",
                     "crosshair",
                     "cover",

codeflash/verification/test_runner.py

@@ -10,7 +10,7 @@
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import custom_addopts, get_run_tmp_file
-from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE, is_compiled_or_bundled_binary
+from codeflash.code_utils.compat import IS_POSIX, get_safe_sys_executable, is_compiled_or_bundled_binary
 from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE
 from codeflash.code_utils.coverage_utils import prepare_coverage_files
 from codeflash.code_utils.shell_utils import get_cross_platform_subprocess_run_args
@@ -98,9 +98,9 @@ def run_behavioral_tests(
                 test_files.append(str(file.instrumented_behavior_file_path))
 
         pytest_cmd_list = (
-            shlex.split(f"{SAFE_SYS_EXECUTABLE} -m pytest", posix=IS_POSIX)
+            shlex.split(f"{get_safe_sys_executable()} -m pytest", posix=IS_POSIX)
             if pytest_cmd == "pytest"
-            else [SAFE_SYS_EXECUTABLE, "-m", *shlex.split(pytest_cmd, posix=IS_POSIX)]
+            else [get_safe_sys_executable(), "-m", *shlex.split(pytest_cmd, posix=IS_POSIX)]
         )
         test_files = list(set(test_files))  # remove multiple calls in the same test function
 
@@ -147,11 +147,14 @@ def run_behavioral_tests(
                         coverage_database_file.unlink()
             else:
                 cov_erase = execute_test_subprocess(
-                    shlex.split(f"{SAFE_SYS_EXECUTABLE} -m coverage erase"), cwd=cwd, env=pytest_test_env, timeout=30
+                    shlex.split(f"{get_safe_sys_executable()} -m coverage erase"),
+                    cwd=cwd,
+                    env=pytest_test_env,
+                    timeout=30,
                 )  # this cleanup is necessary to avoid coverage data from previous runs, if there are any, then the current run will be appended to the previous data, which skews the results
                 logger.debug(cov_erase)
             coverage_cmd = [
-                SAFE_SYS_EXECUTABLE,
+                get_safe_sys_executable(),
                 "-m",
                 "coverage",
                 "run",
@@ -214,7 +217,7 @@ def run_line_profile_tests(
 ) -> tuple[Path, subprocess.CompletedProcess]:
     if test_framework in {"pytest", "unittest"}:  # pytest runs both pytest and unittest tests
         pytest_cmd_list = (
-            shlex.split(f"{SAFE_SYS_EXECUTABLE} -m pytest", posix=IS_POSIX)
+            shlex.split(f"{get_safe_sys_executable()} -m pytest", posix=IS_POSIX)
             if pytest_cmd == "pytest"
             else shlex.split(pytest_cmd)
         )
@@ -276,7 +279,7 @@ def run_benchmarking_tests(
 ) -> tuple[Path, subprocess.CompletedProcess]:
     if test_framework in {"pytest", "unittest"}:  # pytest runs both pytest and unittest tests
         pytest_cmd_list = (
-            shlex.split(f"{SAFE_SYS_EXECUTABLE} -m pytest", posix=IS_POSIX)
+            shlex.split(f"{get_safe_sys_executable()} -m pytest", posix=IS_POSIX)
             if pytest_cmd == "pytest"
             else shlex.split(pytest_cmd)
         )

tests/test_codeflash_capture.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 from codeflash.code_utils.code_utils import get_run_tmp_file
-from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
+from codeflash.code_utils.compat import get_safe_sys_executable
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import FunctionParent, TestFile, TestFiles, TestingMode, TestType, VerificationType
 from codeflash.optimization.function_optimizer import FunctionOptimizer
@@ -54,7 +54,7 @@ def __init__(self):
         with sample_code_path.open("w") as f:
             f.write(sample_code)
         result = execute_test_subprocess(
-            cwd=test_dir, cmd_list=[f"{SAFE_SYS_EXECUTABLE}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
+            cwd=test_dir, cmd_list=[f"{get_safe_sys_executable()}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
         )
         assert not result.stderr
         assert result.returncode == 0
@@ -129,7 +129,7 @@ def __init__(self):
         with sample_code_path.open("w") as f:
             f.write(sample_code)
         result = execute_test_subprocess(
-            cwd=test_dir, cmd_list=[f"{SAFE_SYS_EXECUTABLE}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
+            cwd=test_dir, cmd_list=[f"{get_safe_sys_executable()}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
         )
         assert not result.stderr
         assert result.returncode == 0
@@ -194,7 +194,7 @@ def __init__(self):
         with sample_code_path.open("w") as f:
             f.write(sample_code)
         result = execute_test_subprocess(
-            cwd=test_dir, cmd_list=[f"{SAFE_SYS_EXECUTABLE}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
+            cwd=test_dir, cmd_list=[f"{get_safe_sys_executable()}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
         )
         assert not result.stderr
         assert result.returncode == 0
@@ -279,7 +279,7 @@ def __init__(self):
 
         # Run pytest as a subprocess
         result = execute_test_subprocess(
-            cwd=test_dir, cmd_list=[f"{SAFE_SYS_EXECUTABLE}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
+            cwd=test_dir, cmd_list=[f"{get_safe_sys_executable()}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
         )
 
         # Check for errors
@@ -356,7 +356,7 @@ def __init__(self):
         with sample_code_path.open("w") as f:
             f.write(sample_code)
         result = execute_test_subprocess(
-            cwd=test_dir, cmd_list=[f"{SAFE_SYS_EXECUTABLE}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
+            cwd=test_dir, cmd_list=[f"{get_safe_sys_executable()}", "-m", "pytest", test_file_name, "-s"], env=os.environ.copy()
         )
         assert not result.stderr
         assert result.returncode == 0
@@ -1231,7 +1231,7 @@ def __init__(self):
         test_env.pop("CODEFLASH_TEST_CLASS", None)
 
         result = execute_test_subprocess(
-            cwd=test_dir, cmd_list=[f"{SAFE_SYS_EXECUTABLE}", "-m", "pytest", test_file_name, "-s"], env=test_env
+            cwd=test_dir, cmd_list=[f"{get_safe_sys_executable()}", "-m", "pytest", test_file_name, "-s"], env=test_env
         )
         assert result.returncode == 0
         pattern = r"TEST_INFO_START\|\((.*?)\)\|TEST_INFO_END"
@@ -1292,7 +1292,7 @@ def __init__(self):
 
         test_env = os.environ.copy()
         result = execute_test_subprocess(
-            cwd=test_dir, cmd_list=[f"{SAFE_SYS_EXECUTABLE}", "-m", "pytest", test_file_name, "-s"], env=test_env
+            cwd=test_dir, cmd_list=[f"{get_safe_sys_executable()}", "-m", "pytest", test_file_name, "-s"], env=test_env
         )
         assert result.returncode == 0
         pattern = r"TEST_INFO_START\|\((.*?)\)\|TEST_INFO_END"