From db1ed0859d775534bf312683939da02b2f693599 Mon Sep 17 00:00:00 2001 From: mohammed Date: Thu, 26 Jun 2025 21:49:48 +0300 Subject: [PATCH 1/3] lsp: discover existing tests and return the optimization to the client --- codeflash/code_utils/formatter.py | 3 +++ codeflash/lsp/beta.py | 20 ++++++++++++++------ 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/codeflash/code_utils/formatter.py b/codeflash/code_utils/formatter.py index f4f4c563e..cd8c2efb3 100644 --- a/codeflash/code_utils/formatter.py +++ b/codeflash/code_utils/formatter.py @@ -109,6 +109,9 @@ def format_code( print_status: bool = True, # noqa exit_on_failure: bool = True, # noqa ) -> str: + if console.quiet: + # lsp mode + exit_on_failure = False with tempfile.TemporaryDirectory() as test_dir_str: if isinstance(path, str): path = Path(path) diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index 797e33115..4249d69e9 100644 --- a/codeflash/lsp/beta.py +++ b/codeflash/lsp/beta.py @@ -1,5 +1,8 @@ from __future__ import annotations +import contextlib +import json +import os from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING @@ -60,13 +63,16 @@ def initialize_function_optimization( @server.feature("discoverFunctionTests") def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]: - current_function = server.optimizer.current_function_being_optimized + fto = server.optimizer.current_function_being_optimized + optimizable_funcs = {fto.file_path: [fto]} + + devnull_writer = Path.open(os.devnull, "w") + with contextlib.redirect_stdout(devnull_writer): + function_to_tests, num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs) - optimizable_funcs = {current_function.file_path: [current_function]} + server.optimizer.discovered_tests = function_to_tests - function_to_tests, num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs) - # mocking in order to get things going - return {"functionName": params.functionName, "status": "success", "generated_tests": str(num_discovered_tests)} + return {"functionName": params.functionName, "status": "success", "discovered_tests": num_discovered_tests} @server.feature("prepareOptimization") @@ -145,6 +151,7 @@ def perform_function_optimization( function_to_optimize_source_code=validated_original_code[current_function.file_path].source_code, original_module_ast=original_module_ast, original_module_path=current_function.file_path, + function_to_tests=server.optimizer.discovered_tests or {}, ) server.optimizer.current_function_optimizer = function_optimizer @@ -214,13 +221,14 @@ def perform_function_optimization( "message": f"No best optimizations found for function {function_to_optimize_qualified_name}", } - optimized_source = best_optimization.candidate.source_code # noqa: F841 + optimized_source = best_optimization.candidate.source_code return { "functionName": params.functionName, "status": "success", "message": "Optimization completed successfully", "extra": f"Speedup: {original_code_baseline.runtime / best_optimization.runtime:.2f}x faster", + "optimization": json.dumps(optimized_source, indent=None), } From 5e92b48071d55397aa2463ca9e28221d1131cb68 Mon Sep 17 00:00:00 2001 From: mohammed Date: Thu, 26 Jun 2025 21:56:37 +0300 Subject: [PATCH 2/3] correctly open dev/null --- codeflash/lsp/beta.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index 4249d69e9..c76df5376 100644 --- a/codeflash/lsp/beta.py +++ b/codeflash/lsp/beta.py @@ -66,7 +66,7 @@ def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOpt fto = server.optimizer.current_function_being_optimized optimizable_funcs = {fto.file_path: [fto]} - devnull_writer = Path.open(os.devnull, "w") + devnull_writer = open(os.devnull, "w") # noqa with contextlib.redirect_stdout(devnull_writer): function_to_tests, num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs) From 361c2cd978a16ee3b9833888358c62dbbb78108b Mon Sep 17 00:00:00 2001 From: mohammed Date: Thu, 26 Jun 2025 22:54:15 +0300 Subject: [PATCH 3/3] fix initialization: no property named server_info --- codeflash/lsp/beta.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index c76df5376..1e1397bb8 100644 --- a/codeflash/lsp/beta.py +++ b/codeflash/lsp/beta.py @@ -58,7 +58,7 @@ def initialize_function_optimization( return {"functionName": params.functionName, "status": "not found", "args": None} fto = optimizable_funcs.popitem()[1][0] server.optimizer.current_function_being_optimized = fto - return {"functionName": params.functionName, "status": "success", "info": fto.server_info} + return {"functionName": params.functionName, "status": "success"} @server.feature("discoverFunctionTests")