Merge branch 'main' into improve-readme

Author: KRRT7

codeflash/api/cfapi.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 import requests
+import sentry_sdk
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
@@ -194,7 +195,8 @@ def get_blocklisted_functions() -> dict[str, set[str]] | dict[str, Any]:
         req.raise_for_status()
         content: dict[str, list[str]] = req.json()
     except Exception as e:
-        logger.error(f"Error getting blocklisted functions: {e}", exc_info=True)
+        logger.error(f"Error getting blocklisted functions: {e}")
+        sentry_sdk.capture_exception(e)
         return {}
 
     return {Path(k).name: {v.replace("()", "") for v in values} for k, values in content.items()}

codeflash/cli_cmds/cli.py

@@ -128,14 +128,22 @@ def process_pyproject_config(args: Namespace) -> Namespace:
     assert args.tests_root is not None, "--tests-root must be specified"
     assert Path(args.tests_root).is_dir(), f"--tests-root {args.tests_root} must be a valid directory"
 
-    assert not (env_utils.get_pr_number() is not None and not 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:
+        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)
+
+        owner, repo_name = get_repo_owner_and_name(repo)
+
+        require_github_app_or_exit(owner, repo_name)
+
     if hasattr(args, "ignore_paths") and args.ignore_paths is not None:
         normalized_ignore_paths = []
         for path in args.ignore_paths:

codeflash/cli_cmds/cmd_init.py

@@ -67,20 +67,27 @@ def init_codeflash() -> None:
 
         did_add_new_key = prompt_api_key()
 
-        setup_info: SetupInfo = collect_setup_info()
+        if should_modify_pyproject_toml():
 
-        configure_pyproject_toml(setup_info)
+            setup_info: SetupInfo = collect_setup_info()
+
+            configure_pyproject_toml(setup_info)
 
         install_github_app()
 
         install_github_actions(override_formatter_check=True)
 
+        module_string = ""
+        if "setup_info" in locals():
+            module_string = f" you selected ({setup_info.module_root})"
+
+
         click.echo(
             f"{LF}"
             f"⚡️ Codeflash is now set up! You can now run:{LF}"
             f"    codeflash --file <path-to-file> --function <function-name> to optimize a function within a file{LF}"
             f"    codeflash --file <path-to-file> to optimize all functions in a file{LF}"
-            f"    codeflash --all to optimize all functions in all files in the module you selected ({setup_info.module_root}){LF}"
+            f"    codeflash --all to optimize all functions in all files in the module{module_string}{LF}"
             f"-or-{LF}"
             f"    codeflash --help to see all options{LF}"
         )
@@ -116,6 +123,30 @@ def ask_run_end_to_end_test(args: Namespace) -> None:
         bubble_sort_path, bubble_sort_test_path = create_bubble_sort_file_and_test(args)
         run_end_to_end_test(args, bubble_sort_path, bubble_sort_test_path)
 
+def should_modify_pyproject_toml() -> bool:
+    """
+    Check if the current directory contains a valid pyproject.toml file with codeflash config
+    If it does, ask the user if they want to re-configure it.
+    """
+    from rich.prompt import Confirm
+    pyproject_toml_path = Path.cwd() / "pyproject.toml"
+    if not pyproject_toml_path.exists():
+        return True
+    try:
+        config, config_file_path = parse_config_file(pyproject_toml_path)
+    except Exception as e:
+        return True
+
+    if "module_root" not in config or config["module_root"] is None or not Path(config["module_root"]).is_dir():
+        return True
+    if "tests_root" not in config or config["tests_root"] is None or not Path(config["tests_root"]).is_dir():
+        return True
+
+    create_toml = Confirm.ask(
+        f"✅ A valid Codeflash config already exists in this project. Do you want to re-configure it?", default=False, show_default=True
+    )
+    return create_toml
+
 
 def collect_setup_info() -> SetupInfo:
     curdir = Path.cwd()
@@ -362,7 +393,7 @@ def check_for_toml_or_setup_file() -> str | None:
     return cast(str, project_name)
 
 
-def install_github_actions(override_formatter_check: bool=False) -> None:
+def install_github_actions(override_formatter_check: bool = False) -> None:
     try:
         config, config_file_path = parse_config_file(override_formatter_check=override_formatter_check)
 

codeflash/cli_cmds/console.py

@@ -7,7 +7,16 @@
 
 from rich.console import Console
 from rich.logging import RichHandler
-from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
 
 from codeflash.cli_cmds.console_constants import SPINNER_TYPES
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
@@ -22,15 +31,26 @@
 console = Console()
 logging.basicConfig(
     level=logging.INFO,
-    handlers=[RichHandler(rich_tracebacks=True, markup=False, console=console, show_path=False, show_time=False)],
+    handlers=[
+        RichHandler(
+            rich_tracebacks=True,
+            markup=False,
+            console=console,
+            show_path=False,
+            show_time=False,
+        )
+    ],
     format=BARE_LOGGING_FORMAT,
 )
 
 logger = logging.getLogger("rich")
-logging.getLogger('parso').setLevel(logging.WARNING)
+logging.getLogger("parso").setLevel(logging.WARNING)
+
 
 def paneled_text(
-    text: str, panel_args: dict[str, str | bool] | None = None, text_args: dict[str, str] | None = None
+    text: str,
+    panel_args: dict[str, str | bool] | None = None,
+    text_args: dict[str, str] | None = None,
 ) -> None:
     """Print text in a panel."""
     from rich.panel import Panel
@@ -57,7 +77,9 @@ def code_print(code_str: str) -> None:
 
 
 @contextmanager
-def progress_bar(message: str, *, transient: bool = False) -> Generator[TaskID, None, None]:
+def progress_bar(
+    message: str, *, transient: bool = False
+) -> Generator[TaskID, None, None]:
     """Display a progress bar with a spinner and elapsed time."""
     progress = Progress(
         SpinnerColumn(next(spinners)),
@@ -69,3 +91,25 @@ def progress_bar(message: str, *, transient: bool = False) -> Generator[TaskID,
     task = progress.add_task(message, total=None)
     with progress:
         yield task
+
+
+@contextmanager
+def test_files_progress_bar(
+    total: int, description: str
+) -> Generator[tuple[Progress, TaskID], None, None]:
+    """Progress bar for test files."""
+    with Progress(
+        SpinnerColumn(next(spinners)),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(
+            complete_style="cyan",
+            finished_style="green",
+            pulse_style="yellow",
+        ),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TimeRemainingColumn(),
+        transient=True,
+    ) as progress:
+        task_id = progress.add_task(description, total=total)
+        yield progress, task_id

codeflash/code_utils/code_replacer.py

@@ -159,7 +159,7 @@ def replace_functions_in_file(
     source_code: str,
     original_function_names: list[str],
     optimized_code: str,
-    preexisting_objects: set[tuple[str, tuple[FunctionParent,...]]],
+    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
 ) -> str:
     parsed_function_names = []
     for original_function_name in original_function_names:
@@ -195,7 +195,7 @@ def replace_functions_and_add_imports(
     function_names: list[str],
     optimized_code: str,
     module_abspath: Path,
-    preexisting_objects: set[tuple[str, tuple[FunctionParent,...]]],
+    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
     project_root_path: Path,
 ) -> str:
     return add_needed_imports_from_module(
@@ -211,7 +211,7 @@ def replace_function_definitions_in_module(
     function_names: list[str],
     optimized_code: str,
     module_abspath: Path,
-    preexisting_objects: set[tuple[str, tuple[FunctionParent,...]]],
+    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
     project_root_path: Path,
 ) -> bool:
     source_code: str = module_abspath.read_text(encoding="utf8")

codeflash/code_utils/config_parser.py

@@ -31,7 +31,9 @@ def find_pyproject_toml(config_file: Path | None = None) -> Path:
     raise ValueError(msg)
 
 
-def parse_config_file(config_file_path: Path | None = None, override_formatter_check: bool=False) -> tuple[dict[str, Any], Path]:
+def parse_config_file(
+    config_file_path: Path | None = None, override_formatter_check: bool = False
+) -> tuple[dict[str, Any], Path]:
     config_file_path = find_pyproject_toml(config_file_path)
     try:
         with config_file_path.open("rb") as f:
@@ -85,7 +87,7 @@ def parse_config_file(config_file_path: Path | None = None, override_formatter_c
         "In pyproject.toml, Codeflash only supports the 'test-framework' as pytest and unittest."
     )
     if len(config["formatter-cmds"]) > 0:
-        #see if this is happening during Github actions setup
+        #see if this is happening during GitHub actions setup
         if not override_formatter_check:
             assert config["formatter-cmds"][0] != "your-formatter $file", (
                 "The formatter command is not set correctly in pyproject.toml. Please set the "

codeflash/code_utils/instrument_existing_tests.py

@@ -9,8 +9,7 @@
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-from codeflash.models.models import FunctionParent, TestingMode
-from codeflash.verification.test_results import VerificationType
+from codeflash.models.models import FunctionParent, TestingMode, VerificationType
 
 if TYPE_CHECKING:
     from collections.abc import Iterable