UX changes

ux change

UX changes

Author: KRRT7

codeflash/code_utils/code_utils.py

@@ -5,14 +5,15 @@
 import re
 import shutil
 import site
+import sys
 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.cli_cmds.console import logger, paneled_text
 from codeflash.code_utils.config_parser import find_pyproject_toml
 
 ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
@@ -213,3 +214,9 @@ def cleanup_paths(paths: list[Path]) -> None:
 def restore_conftest(path_to_content_map: dict[Path, str]) -> None:
     for path, file_content in path_to_content_map.items():
         path.write_text(file_content, encoding="utf8")
+
+
+def exit_with_message(message: str, *, error_on_exit: bool = False) -> None:
+    paneled_text(message, panel_args={"style": "red"})
+
+    sys.exit(1 if error_on_exit else 0)

codeflash/code_utils/env_utils.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import os
-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.code_utils import exit_with_message
 from codeflash.code_utils.formatter import format_code
 from codeflash.code_utils.shell_utils import read_api_key_from_shell_config
 
@@ -24,11 +24,11 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool =
         try:
             format_code(formatter_cmds, tmp_file, print_status=False)
         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."
+            exit_with_message(
+                "⚠️ 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.",
+                error_on_exit=True,
             )
-            if exit_on_failure:
-                sys.exit(1)
+
     return return_code
 
 

codeflash/discovery/functions_to_optimize.py

@@ -17,6 +17,7 @@
 from codeflash.api.cfapi import get_blocklisted_functions, is_function_being_optimized_again
 from codeflash.cli_cmds.console import DEBUG_MODE, console, logger
 from codeflash.code_utils.code_utils import (
+    exit_with_message,
     is_class_defined_in_file,
     module_name_from_file_path,
     path_belongs_to_site_packages,
@@ -179,8 +180,9 @@ def get_functions_to_optimize(
             if only_get_this_function is not None:
                 split_function = only_get_this_function.split(".")
                 if len(split_function) > 2:
-                    msg = "Function name should be in the format 'function_name' or 'class_name.function_name'"
-                    raise ValueError(msg)
+                    exit_with_message(
+                        "Function name should be in the format 'function_name' or 'class_name.function_name'"
+                    )
                 if len(split_function) == 2:
                     class_name, only_function_name = split_function
                 else:
@@ -193,8 +195,9 @@ def get_functions_to_optimize(
                     ):
                         found_function = fn
                 if found_function is None:
-                    msg = f"Function {only_function_name} not found in file {file} or the function does not have a 'return' statement or is a property"
-                    raise ValueError(msg)
+                    exit_with_message(
+                        f"Function {only_function_name} not found in file {file}\nor the function does not have a 'return' statement or is a property"
+                    )
                 functions[file] = [found_function]
         else:
             logger.info("Finding all functions modified in the current git diff ...")

codeflash/optimization/optimizer.py

@@ -72,7 +72,6 @@ def create_function_optimizer(
     def run(self) -> None:
         from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint
         from codeflash.code_utils.code_replacer import normalize_code, normalize_node
-        from codeflash.code_utils.code_utils import cleanup_paths
         from codeflash.code_utils.static_analysis import (
             analyze_imported_modules,
             get_first_top_level_function_or_method_ast,
@@ -283,14 +282,22 @@ def run(self) -> None:
             if function_optimizer:
                 function_optimizer.cleanup_generated_files()
 
-            if self.test_cfg.concolic_test_root_dir:
-                cleanup_paths([self.test_cfg.concolic_test_root_dir])
+            self.cleanup_temporary_paths()
+
+    def cleanup_temporary_paths(self) -> None:
+        from codeflash.code_utils.code_utils import cleanup_paths
+
+        cleanup_paths([self.test_cfg.concolic_test_root_dir, self.replay_tests_dir])
 
 
 def run_with_args(args: Namespace) -> None:
+    optimizer = None
     try:
         optimizer = Optimizer(args)
         optimizer.run()
     except KeyboardInterrupt:
-        logger.warning("Keyboard interrupt received. Exiting, please wait…")
+        logger.warning("Keyboard interrupt received. Cleaning up and exiting, please wait…")
+        if optimizer:
+            optimizer.cleanup_temporary_paths()
+
         raise SystemExit from None