diff --git a/python_files/python_server.py b/python_files/python_server.py index a4b15f2cbaae3..deda05753c4d4 100644 --- a/python_files/python_server.py +++ b/python_files/python_server.py @@ -13,28 +13,38 @@ USER_GLOBALS = {} -def send_message(msg: str): +def _send_message(msg: str): length_msg = len(msg) STDOUT.buffer.write(f"Content-Length: {length_msg}\r\n\r\n{msg}".encode()) STDOUT.buffer.flush() +def send_message(**kwargs): + _send_message(json.dumps({"jsonrpc": "2.0", **kwargs})) + + def print_log(msg: str): - send_message(json.dumps({"jsonrpc": "2.0", "method": "log", "params": msg})) + send_message(method="log", params=msg) -def send_response(response: str, response_id: int): - send_message(json.dumps({"jsonrpc": "2.0", "id": response_id, "result": response})) +def send_response( + response: str, + response_id: int, + execution_status: bool = True, # noqa: FBT001, FBT002 +): + send_message( + id=response_id, + result={"status": execution_status, "output": response}, + ) def send_request(params: Optional[Union[List, Dict]] = None): request_id = uuid.uuid4().hex if params is None: - send_message(json.dumps({"jsonrpc": "2.0", "id": request_id, "method": "input"})) + send_message(id=request_id, method="input") else: - send_message( - json.dumps({"jsonrpc": "2.0", "id": request_id, "method": "input", "params": params}) - ) + send_message(id=request_id, method="input", params=params) + return request_id @@ -105,13 +115,14 @@ def execute(request, user_globals): original_stdin = sys.stdin try: sys.stdin = str_input - exec_user_input(request["params"], user_globals) + execution_status = exec_user_input(request["params"], user_globals) finally: sys.stdin = original_stdin - send_response(str_output.get_value(), request["id"]) + + send_response(str_output.get_value(), request["id"], execution_status) -def exec_user_input(user_input, user_globals): +def exec_user_input(user_input, user_globals) -> bool: user_input = user_input[0] if isinstance(user_input, list) else user_input try: @@ -119,10 +130,13 @@ def exec_user_input(user_input, user_globals): retval = callable_(user_input, user_globals) if retval is not None: print(retval) + return True except KeyboardInterrupt: print(traceback.format_exc()) + return False except Exception: print(traceback.format_exc()) + return False class CustomIO(io.TextIOWrapper): diff --git a/src/client/repl/pythonServer.ts b/src/client/repl/pythonServer.ts index e125cf5e8d287..ca45ea900bafb 100644 --- a/src/client/repl/pythonServer.ts +++ b/src/client/repl/pythonServer.ts @@ -9,9 +9,13 @@ import { EventName } from '../telemetry/constants'; const SERVER_PATH = path.join(EXTENSION_ROOT_DIR, 'python_files', 'python_server.py'); let serverInstance: PythonServer | undefined; +export interface ExecutionResult { + status: boolean; + output: string; +} export interface PythonServer extends Disposable { - execute(code: string): Promise; + execute(code: string): Promise; interrupt(): void; input(): void; checkValidCommand(code: string): Promise; @@ -52,8 +56,15 @@ class PythonServerImpl implements Disposable { } @captureTelemetry(EventName.EXECUTION_CODE, { scope: 'selection' }, false) - public execute(code: string): Promise { - return this.connection.sendRequest('execute', code); + public async execute(code: string): Promise { + try { + const result = await this.connection.sendRequest('execute', code); + return result as ExecutionResult; + } catch (err) { + const error = err as Error; + traceError(`Error getting response from REPL server:`, error); + } + return undefined; } public interrupt(): void { diff --git a/src/client/repl/replController.ts b/src/client/repl/replController.ts index 7f11b654c54e6..4760edc980365 100644 --- a/src/client/repl/replController.ts +++ b/src/client/repl/replController.ts @@ -22,27 +22,17 @@ export function createReplController( for (const cell of cells) { const exec = controller.createNotebookCellExecution(cell); exec.start(Date.now()); - try { - const result = await server.execute(cell.document.getText()); - if (result !== '') { - exec.replaceOutput([ - new vscode.NotebookCellOutput([vscode.NotebookCellOutputItem.text(result, 'text/plain')]), - ]); - } - exec.end(true); - } catch (err) { - const error = err as Error; + + const result = await server.execute(cell.document.getText()); + + if (result?.output) { exec.replaceOutput([ - new vscode.NotebookCellOutput([ - vscode.NotebookCellOutputItem.error({ - name: error.name, - message: error.message, - stack: error.stack, - }), - ]), + new vscode.NotebookCellOutput([vscode.NotebookCellOutputItem.text(result.output, 'text/plain')]), ]); - exec.end(false); + // TODO: Properly update via NotebookCellOutputItem.error later. } + + exec.end(result?.status); } }; disposables.push(controller);