Merge remote-tracking branch 'origin/main' into line-profiler

Author: aseembits93

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

@@ -393,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: