From eb96141dc8540f2b6838f7fab278a86481575ede Mon Sep 17 00:00:00 2001 From: Anthony Kim <62267334+anthonykim1@users.noreply.github.com> Date: Tue, 21 Nov 2023 14:12:22 -0800 Subject: [PATCH] Use shell integration to denote success/failure (#22487) Resolves: #22486 Use shell integration to denote success/failure in Python REPL launched from VS Code. This would mean having the blue or red decorators based on whether or not user's command succeeded. --- pythonFiles/pythonrc.py | 54 +++++++++++++++++++ pythonFiles/tests/test_shell_integration.py | 48 +++++++++++++++++ src/client/common/terminal/service.ts | 8 ++- .../common/terminals/service.unit.test.ts | 33 ++++++++++++ 4 files changed, 141 insertions(+), 2 deletions(-) create mode 100644 pythonFiles/pythonrc.py create mode 100644 pythonFiles/tests/test_shell_integration.py diff --git a/pythonFiles/pythonrc.py b/pythonFiles/pythonrc.py new file mode 100644 index 000000000000..a3d09ce9c11c --- /dev/null +++ b/pythonFiles/pythonrc.py @@ -0,0 +1,54 @@ +import sys + +original_ps1 = ">>>" + + +class repl_hooks: + def __init__(self): + self.global_exit = None + self.failure_flag = False + self.original_excepthook = sys.excepthook + self.original_displayhook = sys.displayhook + sys.excepthook = self.my_excepthook + sys.displayhook = self.my_displayhook + + def my_displayhook(self, value): + if value is None: + self.failure_flag = False + + self.original_displayhook(value) + + def my_excepthook(self, type, value, traceback): + self.global_exit = value + self.failure_flag = True + + self.original_excepthook(type, value, traceback) + + +class ps1: + hooks = repl_hooks() + sys.excepthook = hooks.my_excepthook + sys.displayhook = hooks.my_displayhook + + # str will get called for every prompt with exit code to show success/failure + def __str__(self): + exit_code = 0 + if self.hooks.failure_flag: + exit_code = 1 + else: + exit_code = 0 + + # Guide following official VS Code doc for shell integration sequence: + # result = "{command_finished}{prompt_started}{prompt}{command_start}{command_executed}".format( + # command_finished="\x1b]633;D;" + str(exit_code) + "0\x07", + # prompt_started="\x1b]633;A\x07", + # prompt=original_ps1, + # command_start="\x1b]633;B\x07", + # command_executed="\x1b]633;C\x07", + # ) + result = f"{chr(27)}]633;D;{exit_code}0{chr(7)}{chr(27)}]633;A{chr(7)}{original_ps1}{chr(27)}]633;B{chr(7)}{chr(27)}]633;C{chr(7)}" + + return result + + +sys.ps1 = ps1() diff --git a/pythonFiles/tests/test_shell_integration.py b/pythonFiles/tests/test_shell_integration.py new file mode 100644 index 000000000000..e45aeb60470e --- /dev/null +++ b/pythonFiles/tests/test_shell_integration.py @@ -0,0 +1,48 @@ +import importlib +from unittest.mock import Mock + +import pythonrc + + +def test_decoration_success(): + importlib.reload(pythonrc) + ps1 = pythonrc.ps1() + + ps1.hooks.failure_flag = False + result = str(ps1) + assert result == "\x1b]633;D;00\x07\x1b]633;A\x07>>>\x1b]633;B\x07\x1b]633;C\x07" + + +def test_decoration_failure(): + importlib.reload(pythonrc) + ps1 = pythonrc.ps1() + + ps1.hooks.failure_flag = True + result = str(ps1) + + assert result == "\x1b]633;D;10\x07\x1b]633;A\x07>>>\x1b]633;B\x07\x1b]633;C\x07" + + +def test_displayhook_call(): + importlib.reload(pythonrc) + pythonrc.ps1() + mock_displayhook = Mock() + + hooks = pythonrc.repl_hooks() + hooks.original_displayhook = mock_displayhook + + hooks.my_displayhook("mock_value") + + mock_displayhook.assert_called_once_with("mock_value") + + +def test_excepthook_call(): + importlib.reload(pythonrc) + pythonrc.ps1() + mock_excepthook = Mock() + + hooks = pythonrc.repl_hooks() + hooks.original_excepthook = mock_excepthook + + hooks.my_excepthook("mock_type", "mock_value", "mock_traceback") + mock_excepthook.assert_called_once_with("mock_type", "mock_value", "mock_traceback") diff --git a/src/client/common/terminal/service.ts b/src/client/common/terminal/service.ts index 7128d27802f8..cf63603a2e9b 100644 --- a/src/client/common/terminal/service.ts +++ b/src/client/common/terminal/service.ts @@ -2,6 +2,7 @@ // Licensed under the MIT License. import { inject, injectable } from 'inversify'; +import * as path from 'path'; import { CancellationToken, Disposable, Event, EventEmitter, Terminal } from 'vscode'; import '../../common/extensions'; import { IInterpreterService } from '../../interpreter/contracts'; @@ -10,6 +11,8 @@ import { captureTelemetry } from '../../telemetry'; import { EventName } from '../../telemetry/constants'; import { ITerminalAutoActivation } from '../../terminals/types'; import { ITerminalManager } from '../application/types'; +import { EXTENSION_ROOT_DIR } from '../constants'; +import { _SCRIPTS_DIR } from '../process/internal/scripts/constants'; import { IConfigurationService, IDisposableRegistry } from '../types'; import { ITerminalActivator, @@ -28,6 +31,7 @@ export class TerminalService implements ITerminalService, Disposable { private terminalHelper: ITerminalHelper; private terminalActivator: ITerminalActivator; private terminalAutoActivator: ITerminalAutoActivation; + private readonly envVarScript = path.join(EXTENSION_ROOT_DIR, 'pythonFiles', 'pythonrc.py'); public get onDidCloseTerminal(): Event { return this.terminalClosed.event.bind(this.terminalClosed); } @@ -69,14 +73,14 @@ export class TerminalService implements ITerminalService, Disposable { this.terminal!.show(preserveFocus); } } - private async ensureTerminal(preserveFocus: boolean = true): Promise { + public async ensureTerminal(preserveFocus: boolean = true): Promise { if (this.terminal) { return; } this.terminalShellType = this.terminalHelper.identifyTerminalShell(this.terminal); this.terminal = this.terminalManager.createTerminal({ name: this.options?.title || 'Python', - env: this.options?.env, + env: { PYTHONSTARTUP: this.envVarScript }, hideFromUser: this.options?.hideFromUser, }); this.terminalAutoActivator.disableAutoActivation(this.terminal); diff --git a/src/test/common/terminals/service.unit.test.ts b/src/test/common/terminals/service.unit.test.ts index 2f0d86f4000f..7336f7094e6e 100644 --- a/src/test/common/terminals/service.unit.test.ts +++ b/src/test/common/terminals/service.unit.test.ts @@ -2,9 +2,11 @@ // Licensed under the MIT License. import { expect } from 'chai'; +import * as path from 'path'; import * as TypeMoq from 'typemoq'; import { Disposable, Terminal as VSCodeTerminal, WorkspaceConfiguration } from 'vscode'; import { ITerminalManager, IWorkspaceService } from '../../../client/common/application/types'; +import { EXTENSION_ROOT_DIR } from '../../../client/common/constants'; import { IPlatformService } from '../../../client/common/platform/types'; import { TerminalService } from '../../../client/common/terminal/service'; import { ITerminalActivator, ITerminalHelper, TerminalShellType } from '../../../client/common/terminal/types'; @@ -158,6 +160,37 @@ suite('Terminal Service', () => { terminal.verify((t) => t.show(TypeMoq.It.isValue(false)), TypeMoq.Times.exactly(2)); }); + test('Ensure PYTHONSTARTUP is injected', async () => { + service = new TerminalService(mockServiceContainer.object); + terminalActivator + .setup((h) => h.activateEnvironmentInTerminal(TypeMoq.It.isAny(), TypeMoq.It.isAny())) + .returns(() => Promise.resolve(true)) + .verifiable(TypeMoq.Times.once()); + terminalManager + .setup((t) => t.createTerminal(TypeMoq.It.isAny())) + .returns(() => terminal.object) + .verifiable(TypeMoq.Times.atLeastOnce()); + const envVarScript = path.join(EXTENSION_ROOT_DIR, 'pythonFiles', 'pythonrc.py'); + terminalManager + .setup((t) => + t.createTerminal({ + name: TypeMoq.It.isAny(), + env: TypeMoq.It.isObjectWith({ PYTHONSTARTUP: envVarScript }), + hideFromUser: TypeMoq.It.isAny(), + }), + ) + .returns(() => terminal.object) + .verifiable(TypeMoq.Times.atLeastOnce()); + await service.show(); + await service.show(); + await service.show(); + await service.show(); + + terminalHelper.verifyAll(); + terminalActivator.verifyAll(); + terminal.verify((t) => t.show(TypeMoq.It.isValue(true)), TypeMoq.Times.atLeastOnce()); + }); + test('Ensure terminal is activated once after creation', async () => { service = new TerminalService(mockServiceContainer.object); terminalActivator