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.

Author: anthonykim1

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()

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")

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<void> {
         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<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);

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