passes all tests, minor modifications to vsc ext code

Author: aseembits93

codeflash/code_utils/edit_generated_tests.py

@@ -55,14 +55,14 @@ def __init__(self, qualifed_name: str, source_code: str) -> None:
         self.name = qualifed_name.split(".")[-1]
         self.results: list[int] = []  # map actual line number to line number in ast
 
-    def visit_Call(self, node):  # noqa: ANN201, ANN001
+    def visit_Call(self, node):  # type: ignore[no-untyped-def] # noqa: ANN201, ANN001
         """Detect fn calls."""
-        func_name = self._get_called_func_name(node.func)
+        func_name = self._get_called_func_name(node.func)  # type: ignore[no-untyped-call]
         if func_name == self.name:
             self.results.append(node.lineno - 1)
         self.generic_visit(node)
 
-    def _get_called_func_name(self, node):  # noqa: ANN001, ANN202
+    def _get_called_func_name(self, node):  # type: ignore[no-untyped-def] # noqa: ANN001, ANN202
         """Return name of called fn."""
         if isinstance(node, ast.Name):
             return node.id
@@ -130,10 +130,12 @@ def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.Fu
             return updated_node
 
         def leave_SimpleStatementLine(
-            self, original_node: cst.SimpleStatementLine, updated_node: cst.SimpleStatementLine
+            self,
+            original_node: cst.SimpleStatementLine,  # noqa: ARG002
+            updated_node: cst.SimpleStatementLine,
         ) -> cst.SimpleStatementLine:
             # Check if this statement line contains a call to self.name
-            if self._contains_myfunc_call(updated_node):
+            if self._contains_myfunc_call(updated_node):  # type: ignore[no-untyped-call]
                 # Find matching test cases by looking for this test function name in the test results
                 self.cfo_idx_loc_to_look_at += 1
                 matching_original_times = []
@@ -209,21 +211,21 @@ def leave_SimpleStatementLine(
                         )
             return updated_node
 
-        def _contains_myfunc_call(self, node):
+        def _contains_myfunc_call(self, node):  # type: ignore[no-untyped-def] # noqa : ANN202, ANN001
             """Recursively search for any Call node in the statement whose function is named self.name (including obj.myfunc)."""
 
             class Finder(cst.CSTVisitor):
-                def __init__(self, name: str):
+                def __init__(self, name: str) -> None:
                     super().__init__()
                     self.found = False
                     self.name = name
 
-                def visit_Call(self, call_node):
+                def visit_Call(self, call_node) -> None:  # type: ignore[no-untyped-def] # noqa : ANN001
                     func_expr = call_node.func
                     if isinstance(func_expr, cst.Name):
                         if func_expr.value == self.name:
                             self.found = True
-                    elif isinstance(func_expr, cst.Attribute):
+                    elif isinstance(func_expr, cst.Attribute):  # noqa : SIM102
                         if func_expr.attr.value == self.name:
                             self.found = True
 

codeflash/lsp/server.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-from lsprotocol.types import INITIALIZE, MessageType, LogMessageParams
+from lsprotocol.types import INITIALIZE, LogMessageParams, MessageType
 from pygls import uris
 from pygls.protocol import LanguageServerProtocol, lsp_method
 from pygls.server import LanguageServer
@@ -58,21 +58,22 @@ def initialize_optimizer(self, config_file: Path) -> None:
 
     def show_message_log(self, message: str, message_type: str) -> None:
         """Send a log message to the client's output channel.
-        
+
         Args:
             message: The message to log
             message_type: String type - "Info", "Warning", "Error", or "Log"
+
         """
         # Convert string message type to LSP MessageType enum
         type_mapping = {
             "Info": MessageType.Info,
-            "Warning": MessageType.Warning, 
+            "Warning": MessageType.Warning,
             "Error": MessageType.Error,
-            "Log": MessageType.Log
+            "Log": MessageType.Log,
         }
-        
+
         lsp_message_type = type_mapping.get(message_type, MessageType.Info)
-        
+
         # Send log message to client (appears in output channel)
         log_params = LogMessageParams(type=lsp_message_type, message=message)
         self.lsp.notify("window/logMessage", log_params)

codeflash/lsp/server_entry.py

@@ -1,4 +1,5 @@
-"""This script is the dedicated entry point for the Codeflash Language Server.
+"""The following script is the dedicated entry point for the Codeflash Language Server.
+
 It initializes the server and redirects its logs to stderr so that the
 VS Code client can display them in the output channel.
 
@@ -13,7 +14,7 @@
 
 
 # Configure logging to stderr for VS Code output channel
-def setup_logging():
+def setup_logging():  # noqa : ANN201
     # Clear any existing handlers to prevent conflicts
     root_logger = logging.getLogger()
     root_logger.handlers.clear()