Merge pull request #1310 from codeflash-ai/verify-setup

Verify setup issues

cli/code_to_optimize/code_directories/my-best-repo/tests/test_full_bubble_coverage.py

@@ -0,0 +1,47 @@
+import pytest
+from bubble_sort import sorter, sorter_one_level_depth, decompress_braces, sorter_one_level_depth_lower, add_one_level_depth, add
+
+
+def test_sort():
+    input = [5, 4, 3, 2, 1, 0]
+    output = sorter(input)
+    assert output == [0, 1, 2, 3, 4, 5]
+
+    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
+    output = sorter(input)
+    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
+
+    input = list(reversed(range(5000)))
+    output = sorter(input)
+    assert output == list(range(5000))
+def test_sort():
+    input = [5, 4, 3, 2, 1, 0]
+    output = sorter(input)
+    assert output == [0, 1, 2, 3, 4, 5]
+
+    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
+    output = sorter(input)
+    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
+
+    input = list(reversed(range(5000)))
+    output = sorter(input)
+    assert output == list(range(5000))
+
+
+def test_sorter_one_level_depth():
+    input = [3, 2, 1]
+    output = sorter_one_level_depth(input)
+    assert output == [1, 2, 3]
+
+
+def test_add_one_level_depth():
+    assert add_one_level_depth(1, 2) == 3
+    assert add_one_level_depth(-1, 1) == 0
+    assert add_one_level_depth(0, 0) == 0
+
+
+def test_add():
+    assert add(1, 2) == 3
+    assert add(-1, 1) == 0
+    assert add(0, 0) == 0
+

cli/codeflash/api/cfapi.py

@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional
 import requests
 from pydantic.json import pydantic_encoder
 
-from codeflash.cli_cmds.console import logger
+from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
 from codeflash.code_utils.git_utils import get_repo_owner_and_name
 
@@ -21,6 +21,7 @@ if TYPE_CHECKING:
 if os.environ.get("CODEFLASH_CFAPI_SERVER", default="prod").lower() == "local":
     CFAPI_BASE_URL = "http://localhost:3001"
     logger.info(f"Using local CF API at {CFAPI_BASE_URL}.")
+    console.rule()
 else:
     CFAPI_BASE_URL = "https://app.codeflash.ai"
 

cli/codeflash/cli_cmds/cli.py

@@ -123,12 +123,10 @@ def process_pyproject_config(args: Namespace) -> Namespace:
             or not hasattr(args, key.replace("-", "_"))
         ):
             setattr(args, key.replace("-", "_"), pyproject_config[key])
-    assert (
-        args.module_root is not None and Path(args.module_root).is_dir()
-    ), f"--module-root {args.module_root} must be a valid directory"
-    assert (
-        args.tests_root is not None and Path(args.tests_root).is_dir()
-    ), f"--tests-root {args.tests_root} must be a valid directory"
+    assert args.module_root is not None, "--module-root must be specified"
+    assert Path(args.module_root).is_dir(), f"--module-root {args.module_root} must be a valid directory"
+    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 "

cli/codeflash/cli_cmds/cli_common.py

@@ -7,12 +7,16 @@ from typing import Callable, cast
 import click
 import inquirer
 
+from codeflash.cli_cmds.console import console, logger
+
 
 def apologize_and_exit() -> None:
-    click.echo(
+    console.rule()
+    logger.info(
         "💡 If you're having trouble, see https://docs.codeflash.ai/getting-started/local-installation for further help getting started with Codeflash!"
     )
-    click.echo("👋 Exiting...")
+    console.rule()
+    logger.info("👋 Exiting...")
     sys.exit(1)
 
 

cli/codeflash/cli_cmds/cmd_init.py

@@ -19,6 +19,7 @@ from returns.pipeline import is_successful
 
 from codeflash.api.cfapi import is_github_app_installed_on_repo
 from codeflash.cli_cmds.cli_common import apologize_and_exit, inquirer_wrapper, inquirer_wrapper_path
+from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.compat import LF
 from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.code_utils.env_utils import get_codeflash_api_key
@@ -88,11 +89,19 @@ def init_codeflash() -> None:
 
 
 def ask_run_end_to_end_test(args: Namespace) -> None:
-    run_tests = inquirer_wrapper(
-        inquirer.confirm,
-        message="⚡️ Do you want to run a sample optimization to make sure everything's set up correctly? (takes about 3 minutes)",
-        default=True,
+    from rich.prompt import Confirm
+
+    run_tests = Confirm.ask(
+        "⚡️ Do you want to run a sample optimization to make sure everything's set up correctly? (takes about 3 minutes)",
+        choices=["y", "n"],
+        default="y",
+        show_choices=True,
+        show_default=False,
+        console=console,
     )
+
+    console.rule()
+
     if run_tests:
         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)
@@ -617,38 +626,58 @@ def test_sort():
     output = sorter(input)
     assert output == list(range(500))
 """
-    else:
-        click.echo(f"❌ Unsupported test framework: {args.test_framework}")
-        apologize_and_exit()
 
     bubble_sort_path = Path(args.module_root) / "bubble_sort.py"
+    if bubble_sort_path.exists():
+        from rich.prompt import Confirm
+
+        overwrite = Confirm.ask(
+            f"🤔 {bubble_sort_path} already exists. Do you want to overwrite it?", default=True, show_default=False
+        )
+        if not overwrite:
+            apologize_and_exit()
+        console.rule()
+
     bubble_sort_path.write_text(bubble_sort_content, encoding="utf8")
 
     bubble_sort_test_path = Path(args.tests_root) / "test_bubble_sort.py"
     bubble_sort_test_path.write_text(bubble_sort_test_content, encoding="utf8")
 
-    click.echo(f"✅ Created {bubble_sort_path}")
-    click.echo(f"✅ Created {bubble_sort_test_path}")
+    for path in [bubble_sort_path, bubble_sort_test_path]:
+        logger.info(f"✅ Created {path}")
+        console.rule()
 
     return str(bubble_sort_path), str(bubble_sort_test_path)
 
 
 def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test_path: str) -> None:
     command = ["codeflash", "--file", "bubble_sort.py", "--function", "sorter"]
-    sys.stdout.write("Running sample optimization…")
-    sys.stdout.flush()
-    try:
-        process = subprocess.run(command, text=True, cwd=args.module_root, check=False)
-    finally:
-        # Delete the bubble_sort.py file after the test
-        Path(bubble_sort_path).unlink(missing_ok=True)
-        Path(bubble_sort_test_path).unlink(missing_ok=True)
-        click.echo(f"{LF}🗑️ Deleted {bubble_sort_path}")
-        click.echo(f"{LF}🗑️ Deleted {bubble_sort_test_path}")
+    logger.info("Running sample optimization…")
+    console.rule()
 
-    if process.returncode == 0:
-        click.echo(f"{LF}✅ End-to-end test passed. Codeflash has been correctly set up!")
-    else:
-        click.echo(
-            f"{LF}❌ End-to-end test failed. Please check the logs above, and take a look at https://docs.codeflash.ai/getting-started/local-installation for help and troubleshooting."
-        )
+    try:
+        output = []
+        with subprocess.Popen(
+            command, text=True, cwd=args.module_root, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        ) as process:
+            if process.stdout:
+                for line in process.stdout:
+                    stripped = line.strip()
+                    console.print(stripped)
+                    output.append(stripped)
+            process.wait()
+        console.rule()
+        if process.returncode == 0:
+            logger.info("End-to-end test passed. Codeflash has been correctly set up!")
+        else:
+            logger.error(
+                "End-to-end test failed. Please check the logs above, and take a look at https://docs.codeflash.ai/getting-started/local-installation for help and troubleshooting."
+            )
+    finally:
+        console.rule()
+        # Delete the bubble_sort.py file after the test
+        logger.info("🧹 Cleaning up…")
+        for path in [bubble_sort_path, bubble_sort_test_path]:
+            console.rule()
+            Path(path).unlink(missing_ok=True)
+            logger.info(f"🗑️  Deleted {path}")

cli/codeflash/code_utils/formatter.py

@@ -3,36 +3,48 @@ from __future__ import annotations
 import os
 import shlex
 import subprocess
-from pathlib import Path
+import sys
+from typing import TYPE_CHECKING
 
 import isort
 
-from codeflash.cli_cmds.console import logger
+from codeflash.cli_cmds.console import console, logger
+
+if TYPE_CHECKING:
+    from pathlib import Path
 
 
 def format_code(formatter_cmds: list[str], path: Path) -> str:
     # TODO: Only allow a particular whitelist of formatters here to prevent arbitrary code execution
+    formatter_name = formatter_cmds[0].lower()
     if not path.exists():
-        raise FileNotFoundError(f"File {path} does not exist. Cannot format the file.")
-    if formatter_cmds[0].lower() == "disabled":
-        new_code = path.read_text(encoding="utf8")
-        return new_code
-    file_token = "$file"
+        msg = f"File {path} does not exist. Cannot format the file."
+        raise FileNotFoundError(msg)
+    if formatter_name == "disabled":
+        return path.read_text(encoding="utf8")
+    file_token = "$file"  # noqa: S105
 
     for command in formatter_cmds:
         formatter_cmd_list = shlex.split(command, posix=os.name != "nt")
         formatter_cmd_list = [str(path) if chunk == file_token else chunk for chunk in formatter_cmd_list]
-        logger.info(f"Formatting code with {' '.join(formatter_cmd_list)} ...")
-
+        console.rule(f"Formatting code with {' '.join(formatter_cmd_list)} ...")
         try:
-            result = subprocess.run(formatter_cmd_list, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False)
+            result = subprocess.run(formatter_cmd_list, capture_output=True, check=False)
             if result.returncode == 0:
                 logger.info("FORMATTING OK")
             else:
                 logger.error(f"Failed to format code with {' '.join(formatter_cmd_list)}")
-        except Exception as e:
-            logger.exception(f"Failed to format code with {' '.join(formatter_cmd_list)}: {e}")
-            # Fall back to original code if formatter fails
+        except FileNotFoundError as e:
+            from rich.panel import Panel
+            from rich.text import Text
+
+            panel = Panel(
+                Text.from_markup(f"⚠️  Formatter command not found: {' '.join(formatter_cmd_list)}", style="bold red"),
+                expand=False,
+            )
+            console.print(panel)
+
+            raise e from None
 
     return path.read_text(encoding="utf8")
 

cli/codeflash/discovery/functions_to_optimize.py

@@ -14,7 +14,7 @@ import libcst as cst
 from pydantic.dataclasses import dataclass
 
 from codeflash.api.cfapi import get_blocklisted_functions
-from codeflash.cli_cmds.console import logger
+from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.code_utils import (
     is_class_defined_in_file,
     module_name_from_file_path,
@@ -168,6 +168,7 @@ def get_functions_to_optimize(
         )
     elif file is not None:
         logger.info("Finding all functions in the file '%s'…", file)
+        console.rule()
         functions = find_all_functions_in_file(file)
         if only_get_this_function is not None:
             split_function = only_get_this_function.split(".")

cli/codeflash/optimization/optimizer.py

@@ -131,14 +131,17 @@ class Optimizer:
 
             console.rule()
             logger.info(f"Discovering existing unit tests in {self.test_cfg.tests_root}…")
+            console.rule()
             function_to_tests: dict[str, list[FunctionCalledInTest]] = discover_unit_tests(self.test_cfg)
             num_discovered_tests: int = sum([len(value) for value in function_to_tests.values()])
+            console.rule()
             logger.info(f"Discovered {num_discovered_tests} existing unit tests in {self.test_cfg.tests_root}")
             console.rule()
             ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests})
 
             for original_module_path in file_to_funcs_to_optimize:
                 logger.info(f"Examining file {original_module_path!s}…")
+                console.rule()
 
                 original_module_code: str = original_module_path.read_text(encoding="utf8")
                 try:
@@ -731,6 +734,7 @@ class Optimizer:
         func_qualname = function_to_optimize.qualified_name_with_modules_from_root(self.args.project_root)
         if func_qualname not in function_to_tests:
             logger.info(f"Did not find any pre-existing tests for '{func_qualname}', will only use generated tests.")
+            console.rule()
         else:
             test_file_invocation_positions = defaultdict(list)
             for tests_in_file in function_to_tests.get(func_qualname):
@@ -1105,9 +1109,11 @@ class Optimizer:
 
             if compare_test_results(initial_loop_original_test_results, initial_loop_candidate_results):
                 logger.info("Test results matched!")
+                console.rule()
                 equal_results = True
             else:
                 logger.info("Test results did not match the test results of the original code.")
+                console.rule()
                 success = False
                 equal_results = False
 

cli/tests/test_formatter.py

@@ -197,6 +197,5 @@ def foo():
         tmp.write(original_code)
         tmp.flush()
         tmp_path = tmp.name
-
-        actual = format_code(formatter_cmds=["exit 1"], path=Path(tmp_path))
-        assert actual == expected
+        with pytest.raises(FileNotFoundError):
+            format_code(formatter_cmds=["exit 1"], path=Path(tmp_path))