fixes for cli console logs

Author: mohammedahmed18

codeflash/cli_cmds/console.py

@@ -52,7 +52,7 @@
     setattr(
         logger,
         level,
-        lambda msg, _real_fn=real_fn, _level=level, *args, **kwargs: enhanced_log(
+        lambda msg, *args, _real_fn=real_fn, _level=level, **kwargs: enhanced_log(
             msg, _real_fn, _level, *args, **kwargs
         ),
     )

codeflash/lsp/lsp_logger.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.lsp.lsp_message import LspTextMessage
@@ -10,9 +10,9 @@
 @dataclass
 class LspMessageTags:
     # always set default values for message tags
-    not_lsp: bool = False  # !lsp         (prevent the message from being sent to the LSP)
-    force_lsp: bool = False  # lsp          (you can use this to force a message to be sent to the LSP even if the level is not supported)
-    loading: bool = False  # loading      (you can use this to indicate that the message is a loading message)
+    not_lsp: bool = False  # !lsp           (prevent the message from being sent to the LSP)
+    force_lsp: bool = False  # lsp            (you can use this to force a message to be sent to the LSP even if the level is not supported)
+    loading: bool = False  # loading        (you can use this to indicate that the message is a loading message)
 
     h1: bool = False  # h1
     h2: bool = False  # h2
@@ -32,14 +32,15 @@ def add_heading_tags(msg: str, tags: LspMessageTags) -> str:
     return msg
 
 
-def extract_tags(msg: str) -> tuple[LspMessageTags, str]:
+def extract_tags(msg: str) -> tuple[Optional[LspMessageTags], str]:
     if not isinstance(msg, str):
-        return LspMessageTags(), msg
+        return None, msg
 
     parts = msg.split("|tags|")
     if len(parts) == 2:
         message_tags = LspMessageTags()
-        tags = [tag.strip() for tag in parts[0].split(",")]
+        # Use set for O(1) lookups and remove whitespace in a single pass
+        tags = {tag.strip() for tag in parts[0].split(",")}
         if "!lsp" in tags:
             message_tags.not_lsp = True
         if "lsp" in tags:
@@ -56,7 +57,7 @@ def extract_tags(msg: str) -> tuple[LspMessageTags, str]:
             message_tags.h4 = True
         return message_tags, parts[1]
 
-    return LspMessageTags(), msg
+    return None, msg
 
 
 supported_lsp_log_levels = ("info", "debug")
@@ -70,31 +71,23 @@ def enhanced_log(
     **kwargs: Any,  # noqa: ANN401
 ) -> None:
     lsp_enabled = is_LSP_enabled()
-    if not lsp_enabled or not isinstance(msg, str):
-        actual_log_fn(msg, *args, **kwargs)
-        return
-
-    is_lsp_json_message = msg.startswith('{"type"')
-    is_normal_text_message = not is_lsp_json_message
-
-    tags = LspMessageTags()
-    clean_msg = msg
+    tags, clean_msg = extract_tags(msg)
 
-    if is_normal_text_message:
-        tags, clean_msg = extract_tags(msg)
-
-    # normal cli mode
-    if not lsp_enabled:
+    if not lsp_enabled or not isinstance(clean_msg, str):
         actual_log_fn(clean_msg, *args, **kwargs)
         return
 
     #### LSP mode ####
+    is_lsp_json_message = clean_msg.startswith('{"type"')
+    is_normal_text_message = not is_lsp_json_message
+    final_tags = tags if tags else LspMessageTags()
+
     unsupported_level = level not in supported_lsp_log_levels
-    if not tags.force_lsp and (tags.not_lsp or unsupported_level):
+    if not final_tags.force_lsp and (final_tags.not_lsp or unsupported_level):
         return
 
     if is_normal_text_message:
-        clean_msg = add_heading_tags(clean_msg, tags)
-        clean_msg = LspTextMessage(text=clean_msg, takes_time=tags.loading).serialize()
+        clean_msg = add_heading_tags(clean_msg, final_tags)
+        clean_msg = LspTextMessage(text=clean_msg, takes_time=final_tags.loading).serialize()
 
     actual_log_fn(clean_msg, *args, **kwargs)

codeflash/optimization/function_optimizer.py

@@ -319,10 +319,9 @@ def generate_and_instrument_tests(
 
         generated_tests: GeneratedTestsList
         optimizations_set: OptimizationSet
-        generated_tests, function_to_concolic_tests, concolic_test_str, optimizations_set = generated_results.unwrap()
-        count_tests = len(generated_tests.generated_tests)
-        if concolic_test_str:
-            count_tests += 1
+        count_tests, generated_tests, function_to_concolic_tests, concolic_test_str, optimizations_set = (
+            generated_results.unwrap()
+        )
 
         for i, generated_test in enumerate(generated_tests.generated_tests):
             with generated_test.behavior_file_path.open("w", encoding="utf8") as f:
@@ -789,7 +788,7 @@ def log_successful_optimization(
                 f"Performance: {explanation.perf_improvement_line}",
                 "",
                 "#### Explanation\n",
-                explanation.to_console_string(),
+                explanation.__str__(),
             ]
 
             # TODO: display the tests
@@ -807,7 +806,7 @@ def log_successful_optimization(
             explanation_panel = Panel(
                 f"⚡️ Optimization successful! 📄 {self.function_to_optimize.qualified_name} in {explanation.file_path}\n"
                 f"📈 {explanation.perf_improvement_line}\n"
-                f"Explanation: \n{explanation.to_console_string()}",
+                f"Explanation: \n{explanation.__str__()}",
                 title="Optimization Summary",
                 border_style="green",
             )
@@ -1120,10 +1119,16 @@ def generate_tests_and_optimizations(
             logger.warning(f"Failed to generate and instrument tests for {self.function_to_optimize.function_name}")
             return Failure(f"/!\\ NO TESTS GENERATED for {self.function_to_optimize.function_name}")
         function_to_concolic_tests, concolic_test_str = future_concolic_tests.result()
-        logger.info(f"Generated '{len(tests)}' tests for {self.function_to_optimize.function_name}")
+
+        count_tests = len(tests)
+        if concolic_test_str:
+            count_tests += 1
+
+        logger.info(f"Generated '{count_tests}' tests for {self.function_to_optimize.function_name}")
         console.rule()
         generated_tests = GeneratedTestsList(generated_tests=tests)
         result = (
+            count_tests,
             generated_tests,
             function_to_concolic_tests,
             concolic_test_str,
@@ -1622,7 +1627,7 @@ def run_optimized_candidate(
                 console.rule()
                 return Failure("Test results did not match the test results of the original code.")
 
-            logger.info(f"loading|tags|Running performance tests for candidate {optimization_candidate_index}. ⚡️")
+            logger.info(f"loading|tags|Running performance tests ⚡️ for candidate {optimization_candidate_index}.")
 
             if test_framework == "pytest":
                 candidate_benchmarking_results, _ = self.run_and_parse_tests(

codeflash/optimization/optimizer.py

@@ -237,16 +237,16 @@ def discover_tests(
     ) -> tuple[dict[str, set[FunctionCalledInTest]], int]:
         from codeflash.discovery.discover_unit_tests import discover_unit_tests
 
-        logger.info("loading|tags|Scanning project for tests to guide optimization.")
         console.rule()
-        start_time = time.time()
-        function_to_tests, num_discovered_tests, num_discovered_replay_tests = discover_unit_tests(
-            self.test_cfg, file_to_funcs_to_optimize=file_to_funcs_to_optimize
-        )
-        console.rule()
-        logger.info(
-            f"Discovered {num_discovered_tests} existing unit tests and {num_discovered_replay_tests} replay tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"
-        )
+        with progress_bar("Scanning project for tests to guide optimization."):
+            start_time = time.time()
+            function_to_tests, num_discovered_tests, num_discovered_replay_tests = discover_unit_tests(
+                self.test_cfg, file_to_funcs_to_optimize=file_to_funcs_to_optimize
+            )
+            console.rule()
+            logger.info(
+                f"Discovered {num_discovered_tests} existing unit tests and {num_discovered_replay_tests} replay tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"
+            )
         console.rule()
         ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests})
         return function_to_tests, num_discovered_tests

codeflash/result/explanation.py

@@ -41,7 +41,7 @@ def speedup_x(self) -> str:
     def speedup_pct(self) -> str:
         return f"{self.speedup * 100:,.0f}%"
 
-    def to_console_string(self) -> str:
+    def __str__(self) -> str:
         # TODO: After doing the best optimization, remove the test cases that errored on the new code, because they might be failing because of syntax errors and such.
         # TODO: Sometimes the explanation says something similar to "This is the code that was optimized", remove such parts
         original_runtime_human = humanize_runtime(self.original_runtime_ns)