Skip to content

Commit

Permalink
Use shell integration to denote success/failure (#22487)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
anthonykim1 authored Nov 21, 2023
1 parent f6e1338 commit eb96141
Show file tree
Hide file tree
Showing 4 changed files with 141 additions and 2 deletions.
54 changes: 54 additions & 0 deletions pythonFiles/pythonrc.py
Original file line number Diff line number Diff line change
@@ -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()
48 changes: 48 additions & 0 deletions pythonFiles/tests/test_shell_integration.py
Original file line number Diff line number Diff line change
@@ -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")
8 changes: 6 additions & 2 deletions src/client/common/terminal/service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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,
Expand All @@ -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<void> {
return this.terminalClosed.event.bind(this.terminalClosed);
}
Expand Down Expand Up @@ -69,14 +73,14 @@ export class TerminalService implements ITerminalService, Disposable {
this.terminal!.show(preserveFocus);
}
}
private async ensureTerminal(preserveFocus: boolean = true): Promise<void> {
public async ensureTerminal(preserveFocus: boolean = true): Promise<void> {
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);
Expand Down
33 changes: 33 additions & 0 deletions src/test/common/terminals/service.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit eb96141

Please sign in to comment.