override other log methods and log serialized lsp messages

Author: mohammedahmed18

codeflash/cli_cmds/console.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
+import json
 import logging
-import os
 from contextlib import contextmanager
 from itertools import cycle
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -20,7 +20,7 @@
 
 from codeflash.cli_cmds.console_constants import SPINNER_TYPES
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
-from codeflash.lsp.helpers import lsp_log
+from codeflash.lsp.helpers import enhanced_log, is_LSP_enabled
 
 if TYPE_CHECKING:
     from collections.abc import Generator
@@ -31,7 +31,7 @@
 
 console = Console()
 
-if os.getenv("CODEFLASH_LSP"):
+if is_LSP_enabled():
     console.quiet = True
 
 logging.basicConfig(
@@ -43,11 +43,17 @@
 logger = logging.getLogger("rich")
 logging.getLogger("parso").setLevel(logging.WARNING)
 
-real_info_log_fn = logger.info
-logger.info = lambda msg, *args, **kwargs: lsp_log(msg, real_info_log_fn, *args, **kwargs)
+# override the logger to reformat the messages for the lsp
+for level in ("info", "debug", "warning", "error"):
+    real_fn = getattr(logger, level)
+    setattr(logger, level, lambda msg, *args, _real_fn=real_fn, **kwargs: enhanced_log(msg, _real_fn, *args, **kwargs))
 
-real_debug_log_fn = logger.debug
-logger.debug = lambda msg, *args, **kwargs: lsp_log(msg, real_debug_log_fn, *args, **kwargs)
+
+def lsp_log(serializable: dict[str, Any], *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
+    if not is_LSP_enabled():
+        return
+    msg_str = json.dumps(serializable)
+    logger.info(msg_str, *args, **kwargs)
 
 
 def paneled_text(
@@ -69,8 +75,13 @@ def code_print(code_str: str) -> None:
     """Print code with syntax highlighting."""
     from rich.syntax import Syntax
 
+    formatted_code = Syntax(code_str, "python", line_numbers=True, theme="github-dark")
+    # if is_LSP_enabled():
+    #     logger.debug(formatted_code.__str__())
+    #     return
+
     console.rule()
-    console.print(Syntax(code_str, "python", line_numbers=True, theme="github-dark"))
+    console.print(formatted_code)
     console.rule()
 
 

codeflash/lsp/helpers.py

@@ -8,8 +8,20 @@ def is_LSP_enabled() -> bool:
     return os.getenv("CODEFLASH_LSP", default="false").lower() == "true"
 
 
-def lsp_log(msg: str, actual_log_fn: Callable[[str, Any, Any], None], *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
-    if is_LSP_enabled():
-        actual_log_fn(f"::::{msg}", *args, **kwargs)
-    else:
+def enhanced_log(msg: str, actual_log_fn: Callable[[str, Any, Any], None], *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
+    lsp_enabled = is_LSP_enabled()
+
+    # normal cli moded
+    if not lsp_enabled:
         actual_log_fn(msg, *args, **kwargs)
+        return
+
+    #### LSP mode ####
+    if type(msg) != str:  # noqa: E721
+        return
+
+    if msg.startswith("Nonzero return code"):
+        # skip logging the failed tests msg to the client
+        return
+
+    actual_log_fn(f"::::{msg}", *args, **kwargs)