code print over lsp

Author: mohammedahmed18

codeflash/cli_cmds/cli.py

@@ -111,7 +111,7 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
     from codeflash.code_utils.github_utils import require_github_app_or_exit
 
     is_init: bool = args.command.startswith("init") if args.command else False
-    if args.verbose or is_LSP_enabled():
+    if args.verbose:
         logging_config.set_level(logging.DEBUG, echo_setting=not is_init)
     else:
         logging_config.set_level(logging.INFO, echo_setting=not is_init)

codeflash/cli_cmds/console.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-import json
 import logging
 from contextlib import contextmanager
 from itertools import cycle
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Optional
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -21,12 +20,15 @@
 from codeflash.cli_cmds.console_constants import SPINNER_TYPES
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
 from codeflash.lsp.helpers import enhanced_log, is_LSP_enabled
+from codeflash.lsp.lsp_message import LspCodeMessage
 
 if TYPE_CHECKING:
     from collections.abc import Generator
 
     from rich.progress import TaskID
 
+    from codeflash.lsp.lsp_message import LspMessage
+
 DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG
 
 console = Console()
@@ -49,11 +51,11 @@
     setattr(logger, level, lambda msg, *args, _real_fn=real_fn, **kwargs: enhanced_log(msg, _real_fn, *args, **kwargs))
 
 
-def lsp_log(serializable: dict[str, Any], *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
+def lsp_log(message: LspMessage, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
     if not is_LSP_enabled():
         return
-    msg_str = json.dumps(serializable)
-    logger.info(msg_str, *args, **kwargs)
+    json_msg = message.serialize()
+    logger.info(json_msg, *args, **kwargs)
 
 
 def paneled_text(
@@ -71,7 +73,10 @@ def paneled_text(
     console.print(panel)
 
 
-def code_print(code_str: str) -> None:
+def code_print(code_str: str, file_name: Optional[str] = None, function_name: Optional[str] = None) -> None:
+    if is_LSP_enabled():
+        lsp_log(LspCodeMessage(code=code_str, file_name=file_name, function_name=function_name))
+        return
     """Print code with syntax highlighting."""
     from rich.syntax import Syntax
 

codeflash/lsp/lsp_message.py

@@ -30,8 +30,9 @@ def type(self) -> str:
 
     def serialize(self) -> str:
         data = asdict(self)
-        data["type"] = self.type()
-        return json.dumps(data)
+        # Important: keep type as the first key, for making it easy and fast for the client to know if this is a lsp message before parsing it
+        ordered = {"type": self.type(), **data}
+        return json.dumps(ordered)
 
 
 @dataclass
@@ -45,7 +46,7 @@ def type(self) -> str:
 @dataclass
 class LspCodeMessage(LspMessage):
     code: str
-    path: Optional[Path] = None
+    file_name: Optional[Path] = None
     function_name: Optional[str] = None
 
     def type(self) -> str:
@@ -66,3 +67,8 @@ class LspStatsMessage(LspMessage):
 
     def type(self) -> str:
         return "stats"
+
+
+if __name__ == "__main__":
+    msg = LspTextMessage(text="Hello World")
+    print(msg.serialize())

codeflash/optimization/function_optimizer.py

@@ -334,7 +334,7 @@ def generate_and_instrument_tests(
                 )
             )
             logger.info(f"Generated test {i + 1}/{count_tests}:")
-            code_print(generated_test.generated_original_test_source)
+            code_print(generated_test.generated_original_test_source, file_name=f"test_{i + 1}.py")
         if concolic_test_str:
             logger.info(f"Generated test {count_tests}/{count_tests}:")
             code_print(concolic_test_str)
@@ -372,7 +372,11 @@ def optimize_function(self) -> Result[BestOptimization, str]:
 
         should_run_experiment, code_context, original_helper_code = initialization_result.unwrap()
 
-        code_print(code_context.read_writable_code.flat)
+        code_print(
+            code_context.read_writable_code.flat,
+            file_name=self.function_to_optimize.file_path,
+            function_name=self.function_to_optimize.function_name,
+        )
 
         test_setup_result = self.generate_and_instrument_tests(  # also generates optimizations
             code_context, should_run_experiment=should_run_experiment
@@ -496,7 +500,7 @@ def determine_best_candidate(
                 get_run_tmp_file(Path(f"test_return_values_{candidate_index}.bin")).unlink(missing_ok=True)
                 get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite")).unlink(missing_ok=True)
                 logger.info(f"Optimization candidate {candidate_index}/{processor.candidate_len}:")
-                code_print(candidate.source_code.flat)
+                code_print(candidate.source_code.flat, file_name=f"candidate_{candidate_index}.py")
                 # map ast normalized code to diff len, unnormalized code
                 # map opt id to the shortest unnormalized code
                 try: