Merge branch 'main' into saga4/coverage_undefined_error

Author: Saga4

codeflash/api/cfapi.py

@@ -179,9 +179,6 @@ def get_blocklisted_functions() -> dict[str, set[str]] | dict[str, Any]:
     if pr_number is None:
         return {}
 
-    not_found = 404
-    internal_server_error = 500
-
     owner, repo = get_repo_owner_and_name()
     information = {"pr_number": pr_number, "repo_owner": owner, "repo_name": repo}
     try:

codeflash/benchmarking/replay_test.py

@@ -115,8 +115,7 @@ def create_trace_replay_test_code(
             if function_name == "__init__":
                 ret = {class_name_alias}(*args[1:], **kwargs)
             else:
-                instance = args[0] # self
-                ret = instance{method_name}(*args[1:], **kwargs)
+                ret = {class_name_alias}{method_name}(*args, **kwargs)
             """
     )
 

codeflash/cli_cmds/cli.py

@@ -3,21 +3,12 @@
 from argparse import SUPPRESS, ArgumentParser, Namespace
 from pathlib import Path
 
-import git
-
 from codeflash.cli_cmds import logging_config
 from codeflash.cli_cmds.cli_common import apologize_and_exit
 from codeflash.cli_cmds.cmd_init import init_codeflash, install_github_actions
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.config_parser import parse_config_file
-from codeflash.code_utils.git_utils import (
-    check_and_push_branch,
-    check_running_in_git_repo,
-    confirm_proceeding_with_no_git_repo,
-    get_repo_owner_and_name,
-)
-from codeflash.code_utils.github_utils import get_github_secrets_page_url, require_github_app_or_exit
 from codeflash.version import __version__ as version
 
 
@@ -75,6 +66,13 @@ def parse_args() -> Namespace:
 
 
 def process_and_validate_cmd_args(args: Namespace) -> Namespace:
+    from codeflash.code_utils.git_utils import (
+        check_running_in_git_repo,
+        confirm_proceeding_with_no_git_repo,
+        get_repo_owner_and_name,
+    )
+    from codeflash.code_utils.github_utils import require_github_app_or_exit
+
     is_init: bool = args.command.startswith("init") if args.command else False
     if args.verbose:
         logging_config.set_level(logging.DEBUG, echo_setting=not is_init)
@@ -144,21 +142,26 @@ def process_pyproject_config(args: Namespace) -> Namespace:
         assert Path(args.benchmarks_root).resolve().is_relative_to(Path(args.tests_root).resolve()), (
             f"--benchmarks-root {args.benchmarks_root} must be a subdirectory of --tests-root {args.tests_root}"
         )
-    if env_utils.get_pr_number() is not None:
-        assert env_utils.ensure_codeflash_api_key(), (
-            "Codeflash API key not found. When running in a Github Actions Context, provide the "
-            "'CODEFLASH_API_KEY' environment variable as a secret.\n"
-            "You can add a secret by going to your repository's settings page, then clicking 'Secrets' in the left sidebar.\n"
-            "Then, click 'New repository secret' and add your api key with the variable name CODEFLASH_API_KEY.\n"
-            f"Here's a direct link: {get_github_secrets_page_url()}\n"
-            "Exiting..."
-        )
+        if env_utils.get_pr_number() is not None:
+            import git
+
+            from codeflash.code_utils.git_utils import get_repo_owner_and_name
+            from codeflash.code_utils.github_utils import get_github_secrets_page_url, require_github_app_or_exit
+
+            assert env_utils.ensure_codeflash_api_key(), (
+                "Codeflash API key not found. When running in a Github Actions Context, provide the "
+                "'CODEFLASH_API_KEY' environment variable as a secret.\n"
+                "You can add a secret by going to your repository's settings page, then clicking 'Secrets' in the left sidebar.\n"
+                "Then, click 'New repository secret' and add your api key with the variable name CODEFLASH_API_KEY.\n"
+                f"Here's a direct link: {get_github_secrets_page_url()}\n"
+                "Exiting..."
+            )
 
-        repo = git.Repo(search_parent_directories=True)
+            repo = git.Repo(search_parent_directories=True)
 
-        owner, repo_name = get_repo_owner_and_name(repo)
+            owner, repo_name = get_repo_owner_and_name(repo)
 
-        require_github_app_or_exit(owner, repo_name)
+            require_github_app_or_exit(owner, repo_name)
 
     if hasattr(args, "ignore_paths") and args.ignore_paths is not None:
         normalized_ignore_paths = []
@@ -187,6 +190,11 @@ def project_root_from_module_root(module_root: Path, pyproject_file_path: Path)
 
 def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
     if hasattr(args, "all"):
+        import git
+
+        from codeflash.code_utils.git_utils import check_and_push_branch, get_repo_owner_and_name
+        from codeflash.code_utils.github_utils import require_github_app_or_exit
+
         # Ensure that the user can actually open PRs on the repo.
         try:
             git_repo = git.Repo(search_parent_directories=True)

codeflash/cli_cmds/cmd_init.py

@@ -514,7 +514,8 @@ def install_github_actions(override_formatter_check: bool = False) -> None:  # n
         from importlib.resources import files
 
         benchmark_mode = False
-        if "benchmarks_root" in config:
+        benchmarks_root = config.get("benchmarks_root", "").strip()
+        if benchmarks_root and benchmarks_root != "":
             benchmark_mode = inquirer_wrapper(
                 inquirer.confirm,
                 message="⚡️It looks like you've configured a benchmarks_root in your config. Would you like to run the Github action in benchmark mode? "
@@ -720,7 +721,7 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
         )
     elif formatter == "don't use a formatter":
         formatter_cmds.append("disabled")
-    check_formatter_installed(formatter_cmds)
+    check_formatter_installed(formatter_cmds, exit_on_failure=False)
     codeflash_section["formatter-cmds"] = formatter_cmds
     # Add the 'codeflash' section, ensuring 'tool' section exists
     tool_section = pyproject_data.get("tool", tomlkit.table())

codeflash/code_utils/code_extractor.py

@@ -5,7 +5,6 @@
 from typing import TYPE_CHECKING, Optional
 
 import libcst as cst
-import libcst.matchers as m
 from libcst.codemod import CodemodContext
 from libcst.codemod.visitors import AddImportsVisitor, GatherImportsVisitor, RemoveImportsVisitor
 from libcst.helpers import calculate_module_and_package
@@ -248,6 +247,8 @@ class FutureAliasedImportTransformer(cst.CSTTransformer):
     def leave_ImportFrom(
         self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
     ) -> cst.BaseSmallStatement | cst.FlattenSentinel[cst.BaseSmallStatement] | cst.RemovalSentinel:
+        import libcst.matchers as m
+
         if (
             (updated_node_module := updated_node.module)
             and updated_node_module.value == "__future__"

codeflash/code_utils/code_utils.py

@@ -2,13 +2,87 @@
 
 import ast
 import os
+import re
 import shutil
 import site
+from contextlib import contextmanager
 from functools import lru_cache
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
+import tomlkit
+
 from codeflash.cli_cmds.console import logger
+from codeflash.code_utils.config_parser import find_pyproject_toml
+
+
+@contextmanager
+def custom_addopts() -> None:
+    pyproject_file = find_pyproject_toml()
+    original_content = None
+    non_blacklist_plugin_args = ""
+
+    try:
+        # Read original file
+        if pyproject_file.exists():
+            with Path.open(pyproject_file, encoding="utf-8") as f:
+                original_content = f.read()
+                data = tomlkit.parse(original_content)
+            # Backup original addopts
+            original_addopts = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "")
+            # nothing to do if no addopts present
+            if original_addopts != "":
+                original_addopts = [x.strip() for x in original_addopts]
+                non_blacklist_plugin_args = re.sub(r"-n(?: +|=)\S+", "", " ".join(original_addopts)).split(" ")
+                non_blacklist_plugin_args = [x for x in non_blacklist_plugin_args if x != ""]
+                if non_blacklist_plugin_args != original_addopts:
+                    data["tool"]["pytest"]["ini_options"]["addopts"] = non_blacklist_plugin_args
+                    # Write modified file
+                    with Path.open(pyproject_file, "w", encoding="utf-8") as f:
+                        f.write(tomlkit.dumps(data))
+
+        yield
+
+    finally:
+        # Restore original file
+        if (
+            original_content
+            and pyproject_file.exists()
+            and tuple(original_addopts) not in {(), tuple(non_blacklist_plugin_args)}
+        ):
+            with Path.open(pyproject_file, "w", encoding="utf-8") as f:
+                f.write(original_content)
+
+
+@contextmanager
+def add_addopts_to_pyproject() -> None:
+    pyproject_file = find_pyproject_toml()
+    original_content = None
+    try:
+        # Read original file
+        if pyproject_file.exists():
+            with Path.open(pyproject_file, encoding="utf-8") as f:
+                original_content = f.read()
+                data = tomlkit.parse(original_content)
+            data["tool"]["pytest"] = {}
+            data["tool"]["pytest"]["ini_options"] = {}
+            data["tool"]["pytest"]["ini_options"]["addopts"] = [
+                "-n=auto",
+                "-n",
+                "1",
+                "-n 1",
+                "-n      1",
+                "-n      auto",
+            ]
+            with Path.open(pyproject_file, "w", encoding="utf-8") as f:
+                f.write(tomlkit.dumps(data))
+
+        yield
+
+    finally:
+        # Restore original file
+        with Path.open(pyproject_file, "w", encoding="utf-8") as f:
+            f.write(original_content)
 
 
 def encoded_tokens_len(s: str) -> int:

codeflash/code_utils/config_consts.py

@@ -8,3 +8,4 @@
 N_TESTS_TO_GENERATE = 2
 TOTAL_LOOPING_TIME = 10.0  # 10 second candidate benchmarking budget
 COVERAGE_THRESHOLD = 60.0
+MIN_TESTCASE_PASSED_THRESHOLD = 6

codeflash/code_utils/env_utils.py

@@ -1,25 +1,18 @@
 from __future__ import annotations
 
 import os
-import shlex
-import subprocess
+import sys
 import tempfile
 from functools import lru_cache
 from pathlib import Path
 from typing import Optional
 
 from codeflash.cli_cmds.console import logger
+from codeflash.code_utils.formatter import format_code
 from codeflash.code_utils.shell_utils import read_api_key_from_shell_config
 
 
-class FormatterNotFoundError(Exception):
-    """Exception raised when a formatter is not found."""
-
-    def __init__(self, formatter_cmd: str) -> None:
-        super().__init__(f"Formatter command not found: {formatter_cmd}")
-
-
-def check_formatter_installed(formatter_cmds: list[str]) -> bool:
+def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool = True) -> bool:  # noqa
     return_code = True
     if formatter_cmds[0] == "disabled":
         return return_code
@@ -28,22 +21,14 @@ def check_formatter_installed(formatter_cmds: list[str]) -> bool:
         f.write(tmp_code)
         f.flush()
         tmp_file = Path(f.name)
-        file_token = "$file"  # noqa: S105
-        for command in set(formatter_cmds):
-            formatter_cmd_list = shlex.split(command, posix=os.name != "nt")
-            formatter_cmd_list = [tmp_file.as_posix() if chunk == file_token else chunk for chunk in formatter_cmd_list]
-            try:
-                result = subprocess.run(formatter_cmd_list, capture_output=True, check=False)
-            except (FileNotFoundError, NotADirectoryError):
-                return_code = False
-                break
-            if result.returncode:
-                return_code = False
-                break
-    tmp_file.unlink(missing_ok=True)
-    if not return_code:
-        msg = f"Error running formatter command: {command}"
-        raise FormatterNotFoundError(msg)
+        try:
+            format_code(formatter_cmds, tmp_file)
+        except Exception:
+            print(
+                "⚠️ Codeflash requires a code formatter to be installed in your environment, but none was found. Please install a supported formatter, verify the formatter-cmds in your codeflash pyproject.toml config and try again."
+            )
+            if exit_on_failure:
+                sys.exit(1)
     return return_code
 
 

codeflash/code_utils/formatter.py

@@ -22,7 +22,7 @@ def format_code(formatter_cmds: list[str], path: Path) -> str:
     if formatter_name == "disabled":
         return path.read_text(encoding="utf8")
     file_token = "$file"  # noqa: S105
-    for command in set(formatter_cmds):
+    for command in formatter_cmds:
         formatter_cmd_list = shlex.split(command, posix=os.name != "nt")
         formatter_cmd_list = [path.as_posix() if chunk == file_token else chunk for chunk in formatter_cmd_list]
         try:

codeflash/context/code_context_extractor.py

@@ -4,10 +4,9 @@
 from collections import defaultdict
 from itertools import chain
 from pathlib import Path  # noqa: TC003
+from typing import TYPE_CHECKING
 
-import jedi
 import libcst as cst
-from jedi.api.classes import Name  # noqa: TC002
 from libcst import CSTNode  # noqa: TC002
 
 from codeflash.cli_cmds.console import logger
@@ -24,6 +23,9 @@
 )
 from codeflash.optimization.function_context import belongs_to_function_qualified
 
+if TYPE_CHECKING:
+    from jedi.api.classes import Name
+
 
 def get_code_optimization_context(
     function_to_optimize: FunctionToOptimize,
@@ -354,6 +356,8 @@ def extract_code_markdown_context_from_files(
 def get_function_to_optimize_as_function_source(
     function_to_optimize: FunctionToOptimize, project_root_path: Path
 ) -> FunctionSource:
+    import jedi
+
     # Use jedi to find function to optimize
     script = jedi.Script(path=function_to_optimize.file_path, project=jedi.Project(path=project_root_path))
 
@@ -389,6 +393,8 @@ def get_function_to_optimize_as_function_source(
 def get_function_sources_from_jedi(
     file_path_to_qualified_function_names: dict[Path, set[str]], project_root_path: Path
 ) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]:
+    import jedi
+
     file_path_to_function_source = defaultdict(set)
     function_source_list: list[FunctionSource] = []
     for file_path, qualified_function_names in file_path_to_qualified_function_names.items():