Adding PYTHONSTARTUP with shell integration to environment variable c…

…ollection (microsoft#24111) Resolves: microsoft#23930 - setting to opt out of PYTHONSTARTUP injection. --------- Co-authored-by: Courtney Webster <[email protected]>
anthonykim1 · Sep 19, 2024 · 3fcb3fa · 3fcb3fa
1 parent 07a0755
commit 3fcb3fa
Show file tree

Hide file tree

Showing 14 changed files with 91 additions and 21 deletions.
diff --git a/package.json b/package.json
@@ -656,6 +656,12 @@
                     "scope": "resource",
                     "type": "array"
                 },
+                "python.REPL.enableShellIntegration": {
+                    "default": true,
+                    "description": "%python.REPL.enableShellIntegration.description%",
+                    "scope": "resource",
+                    "type": "boolean"
+                },
                 "python.REPL.enableREPLSmartSend": {
                     "default": true,
                     "description": "%python.EnableREPLSmartSend.description%",

diff --git a/package.nls.json b/package.nls.json
@@ -65,6 +65,7 @@
     "python.pixiToolPath.description": "Path to the pixi executable.",
     "python.EnableREPLSmartSend.description": "Toggle Smart Send for the Python REPL. Smart Send enables sending the smallest runnable block of code to the REPL on Shift+Enter and moves the cursor accordingly.",
     "python.REPL.sendToNativeREPL.description": "Toggle to send code to Python REPL instead of the terminal on execution. Turning this on will change the behavior for both Smart Send and Run Selection/Line in the Context Menu.",
+    "python.REPL.enableShellIntegration.description": "Enable Shell Integration for Python Terminal REPL. Shell Integration enhances the terminal experience by allowing command decorations, run recent command, and improving accessibility for Python REPL in the terminal.",
     "python.tensorBoard.logDirectory.description": "Set this setting to your preferred TensorBoard log directory to skip log directory prompt when starting TensorBoard.",
     "python.tensorBoard.logDirectory.markdownDeprecationMessage": "Tensorboard support has been moved to the extension [Tensorboard extension](https://marketplace.visualstudio.com/items?itemName=ms-toolsai.tensorboard). Instead use the setting `tensorBoard.logDirectory`.",
     "python.tensorBoard.logDirectory.deprecationMessage": "Tensorboard support has been moved to the extension Tensorboard extension. Instead use the setting `tensorBoard.logDirectory`.",

diff --git a/src/client/common/types.ts b/src/client/common/types.ts
@@ -202,6 +202,7 @@ export interface ITerminalSettings {
 export interface IREPLSettings {
     readonly enableREPLSmartSend: boolean;
     readonly sendToNativeREPL: boolean;
+    readonly enableShellIntegration: boolean;
 }
 
 export interface IExperiments {

diff --git a/src/client/common/vscodeApis/workspaceApis.ts b/src/client/common/vscodeApis/workspaceApis.ts
@@ -93,3 +93,11 @@ export function isTrusted(): boolean {
 export function onDidGrantWorkspaceTrust(handler: () => void): vscode.Disposable {
     return vscode.workspace.onDidGrantWorkspaceTrust(handler);
 }
+
+export function createDirectory(uri: vscode.Uri): Thenable<void> {
+    return vscode.workspace.fs.createDirectory(uri);
+}
+
+export function copy(source: vscode.Uri, dest: vscode.Uri, options?: { overwrite?: boolean }): Thenable<void> {
+    return vscode.workspace.fs.copy(source, dest, options);
+}
diff --git a/src/client/extensionActivation.ts b/src/client/extensionActivation.ts
@@ -54,6 +54,7 @@ import { DebuggerTypeName } from './debugger/constants';
 import { StopWatch } from './common/utils/stopWatch';
 import { registerReplCommands, registerReplExecuteOnEnter, registerStartNativeReplCommand } from './repl/replCommands';
 import { registerTriggerForTerminalREPL } from './terminals/codeExecution/terminalReplWatcher';
+import { registerPythonStartup } from './terminals/pythonStartup';
 
 export async function activateComponents(
     // `ext` is passed to any extra activation funcs.
@@ -177,6 +178,8 @@ async function activateLegacy(ext: ExtensionState, startupStopWatch: StopWatch):
 
             serviceManager.get<ITerminalAutoActivation>(ITerminalAutoActivation).register();
 
+            await registerPythonStartup(ext.context);
+
             serviceManager.get<ICodeExecutionManager>(ICodeExecutionManager).registerCommands();
 
             disposables.push(new ReplProvider(serviceContainer));

diff --git a/src/client/terminals/envCollectionActivation/service.ts b/src/client/terminals/envCollectionActivation/service.ts
@@ -37,7 +37,11 @@ import { TerminalShellType } from '../../common/terminal/types';
 import { OSType } from '../../common/utils/platform';
 
 import { PythonEnvType } from '../../pythonEnvironments/base/info';
-import { IShellIntegrationService, ITerminalDeactivateService, ITerminalEnvVarCollectionService } from '../types';
+import {
+    IShellIntegrationDetectionService,
+    ITerminalDeactivateService,
+    ITerminalEnvVarCollectionService,
+} from '../types';
 import { ProgressService } from '../../common/application/progressService';
 
 @injectable()
@@ -80,7 +84,8 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
         @inject(IConfigurationService) private readonly configurationService: IConfigurationService,
         @inject(ITerminalDeactivateService) private readonly terminalDeactivateService: ITerminalDeactivateService,
         @inject(IPathUtils) private readonly pathUtils: IPathUtils,
-        @inject(IShellIntegrationService) private readonly shellIntegrationService: IShellIntegrationService,
+        @inject(IShellIntegrationDetectionService)
+        private readonly shellIntegrationDetectionService: IShellIntegrationDetectionService,
         @inject(IEnvironmentVariablesProvider)
         private readonly environmentVariablesProvider: IEnvironmentVariablesProvider,
     ) {
@@ -113,7 +118,7 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
                     this,
                     this.disposables,
                 );
-                this.shellIntegrationService.onDidChangeStatus(
+                this.shellIntegrationDetectionService.onDidChangeStatus(
                     async () => {
                         traceInfo("Shell integration status changed, can confirm it's working.");
                         await this._applyCollection(undefined).ignoreErrors();
@@ -139,7 +144,7 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
                     this.disposables,
                 );
                 const { shell } = this.applicationEnvironment;
-                const isActive = await this.shellIntegrationService.isWorking();
+                const isActive = await this.shellIntegrationDetectionService.isWorking();
                 const shellType = identifyShellFromShellPath(shell);
                 if (!isActive && shellType !== TerminalShellType.commandPrompt) {
                     traceWarn(
@@ -328,7 +333,7 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
                 // PS1 should be set but no PS1 was set.
                 return;
             }
-            const config = await this.shellIntegrationService.isWorking();
+            const config = await this.shellIntegrationDetectionService.isWorking();
             if (!config) {
                 traceVerbose('PS1 is not set when shell integration is disabled.');
                 return;
@@ -395,7 +400,7 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
     }
 
     private async getPrependOptions(): Promise<EnvironmentVariableMutatorOptions> {
-        const isActive = await this.shellIntegrationService.isWorking();
+        const isActive = await this.shellIntegrationDetectionService.isWorking();
         // Ideally we would want to prepend exactly once, either at shell integration or process creation.
         // TODO: Stop prepending altogether once https://github.com/microsoft/vscode/issues/145234 is available.
         return isActive

diff --git a/src/client/terminals/envCollectionActivation/shellIntegrationService.ts b/src/client/terminals/envCollectionActivation/shellIntegrationService.ts
@@ -14,7 +14,7 @@ import { TerminalShellType } from '../../common/terminal/types';
 import { IDisposableRegistry, IPersistentStateFactory } from '../../common/types';
 import { sleep } from '../../common/utils/async';
 import { traceError, traceVerbose } from '../../logging';
-import { IShellIntegrationService } from '../types';
+import { IShellIntegrationDetectionService } from '../types';
 
 /**
  * This is a list of shells which support shell integration:
@@ -33,7 +33,7 @@ export enum isShellIntegrationWorking {
 }
 
 @injectable()
-export class ShellIntegrationService implements IShellIntegrationService {
+export class ShellIntegrationDetectionService implements IShellIntegrationDetectionService {
     private isWorkingForShell = new Set<TerminalShellType>();
 
     private readonly didChange = new EventEmitter<void>();

diff --git a/src/client/terminals/pythonStartup.ts b/src/client/terminals/pythonStartup.ts
@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { ExtensionContext, Uri } from 'vscode';
+import * as path from 'path';
+import { copy, createDirectory, getConfiguration } from '../common/vscodeApis/workspaceApis';
+import { EXTENSION_ROOT_DIR } from '../constants';
+
+export async function registerPythonStartup(context: ExtensionContext): Promise<void> {
+    const config = getConfiguration('python');
+    const pythonrcSetting = config.get<boolean>('REPL.enableShellIntegration');
+
+    if (pythonrcSetting) {
+        const storageUri = context.storageUri || context.globalStorageUri;
+        try {
+            await createDirectory(storageUri);
+        } catch {
+            // already exists, most likely
+        }
+        const destPath = Uri.joinPath(storageUri, 'pythonrc.py');
+        const sourcePath = path.join(EXTENSION_ROOT_DIR, 'python_files', 'pythonrc.py');
+        await copy(Uri.file(sourcePath), destPath, { overwrite: true });
+        context.environmentVariableCollection.replace('PYTHONSTARTUP', destPath.fsPath);
+    } else {
+        context.environmentVariableCollection.delete('PYTHONSTARTUP');
+    }
+}
diff --git a/src/client/terminals/serviceRegistry.ts b/src/client/terminals/serviceRegistry.ts
@@ -12,16 +12,16 @@ import {
     ICodeExecutionHelper,
     ICodeExecutionManager,
     ICodeExecutionService,
-    IShellIntegrationService,
+    IShellIntegrationDetectionService,
     ITerminalAutoActivation,
     ITerminalDeactivateService,
     ITerminalEnvVarCollectionService,
 } from './types';
 import { TerminalEnvVarCollectionService } from './envCollectionActivation/service';
 import { IExtensionActivationService, IExtensionSingleActivationService } from '../activation/types';
 import { TerminalIndicatorPrompt } from './envCollectionActivation/indicatorPrompt';
-import { ShellIntegrationService } from './envCollectionActivation/shellIntegrationService';
 import { TerminalDeactivateService } from './envCollectionActivation/deactivateService';
+import { ShellIntegrationDetectionService } from './envCollectionActivation/shellIntegrationService';
 
 export function registerTypes(serviceManager: IServiceManager): void {
     serviceManager.addSingleton<ICodeExecutionHelper>(ICodeExecutionHelper, CodeExecutionHelper);
@@ -50,6 +50,10 @@ export function registerTypes(serviceManager: IServiceManager): void {
         IExtensionSingleActivationService,
         TerminalIndicatorPrompt,
     );
-    serviceManager.addSingleton<IShellIntegrationService>(IShellIntegrationService, ShellIntegrationService);
+    serviceManager.addSingleton<IShellIntegrationDetectionService>(
+        IShellIntegrationDetectionService,
+        ShellIntegrationDetectionService,
+    );
+
     serviceManager.addBinding(ITerminalEnvVarCollectionService, IExtensionActivationService);
 }
diff --git a/src/client/terminals/types.ts b/src/client/terminals/types.ts
@@ -42,8 +42,8 @@ export interface ITerminalEnvVarCollectionService {
     isTerminalPromptSetCorrectly(resource?: Resource): boolean;
 }
 
-export const IShellIntegrationService = Symbol('IShellIntegrationService');
-export interface IShellIntegrationService {
+export const IShellIntegrationDetectionService = Symbol('IShellIntegrationDetectionService');
+export interface IShellIntegrationDetectionService {
     onDidChangeStatus: Event<void>;
     isWorking(): Promise<boolean>;
 }
@@ -53,3 +53,8 @@ export interface ITerminalDeactivateService {
     initializeScriptParams(shell: string): Promise<void>;
     getScriptLocation(shell: string, resource: Resource): Promise<string | undefined>;
 }
+
+export const IPythonStartupEnvVarService = Symbol('IPythonStartupEnvVarService');
+export interface IPythonStartupEnvVarService {
+    register(): void;
+}
diff --git a/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts b/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts
@@ -37,7 +37,7 @@ import { IInterpreterService } from '../../../client/interpreter/contracts';
 import { PathUtils } from '../../../client/common/platform/pathUtils';
 import { PythonEnvType } from '../../../client/pythonEnvironments/base/info';
 import { PythonEnvironment } from '../../../client/pythonEnvironments/info';
-import { IShellIntegrationService, ITerminalDeactivateService } from '../../../client/terminals/types';
+import { IShellIntegrationDetectionService, ITerminalDeactivateService } from '../../../client/terminals/types';
 import { IEnvironmentVariablesProvider } from '../../../client/common/variables/types';
 
 suite('Terminal Environment Variable Collection Service', () => {
@@ -58,7 +58,7 @@ suite('Terminal Environment Variable Collection Service', () => {
         title: Interpreters.activatingTerminals,
     };
     let configService: IConfigurationService;
-    let shellIntegrationService: IShellIntegrationService;
+    let shellIntegrationService: IShellIntegrationDetectionService;
     const displayPath = 'display/path';
     const customShell = 'powershell';
     const defaultShell = defaultShells[getOSType()];
@@ -76,7 +76,7 @@ suite('Terminal Environment Variable Collection Service', () => {
         context = mock<IExtensionContext>();
         shell = mock<IApplicationShell>();
         const envVarProvider = mock<IEnvironmentVariablesProvider>();
-        shellIntegrationService = mock<IShellIntegrationService>();
+        shellIntegrationService = mock<IShellIntegrationDetectionService>();
         when(shellIntegrationService.isWorking()).thenResolve(true);
         globalCollection = mock<GlobalEnvironmentVariableCollection>();
         collection = mock<EnvironmentVariableCollection>();

diff --git a/src/test/terminals/codeExecution/helper.test.ts b/src/test/terminals/codeExecution/helper.test.ts
@@ -112,7 +112,12 @@ suite('Terminal - Code Execution Helper', () => {
         activeResourceService.setup((a) => a.getActiveResource()).returns(() => resource);
         pythonSettings
             .setup((s) => s.REPL)
-            .returns(() => ({ enableREPLSmartSend: false, REPLSmartSend: false, sendToNativeREPL: false }));
+            .returns(() => ({
+                enableREPLSmartSend: false,
+                REPLSmartSend: false,
+                sendToNativeREPL: false,
+                enableShellIntegration: true,
+            }));
         configurationService.setup((x) => x.getSettings(TypeMoq.It.isAny())).returns(() => pythonSettings.object);
         configurationService
             .setup((c) => c.getSettings(TypeMoq.It.isAny()))

diff --git a/src/test/terminals/codeExecution/smartSend.test.ts b/src/test/terminals/codeExecution/smartSend.test.ts
@@ -109,7 +109,12 @@ suite('REPL - Smart Send', () => {
 
         pythonSettings
             .setup((s) => s.REPL)
-            .returns(() => ({ enableREPLSmartSend: true, REPLSmartSend: true, sendToNativeREPL: false }));
+            .returns(() => ({
+                enableREPLSmartSend: true,
+                REPLSmartSend: true,
+                sendToNativeREPL: false,
+                enableShellIntegration: true,
+            }));
 
         configurationService.setup((x) => x.getSettings(TypeMoq.It.isAny())).returns(() => pythonSettings.object);
 

diff --git a/src/test/terminals/serviceRegistry.unit.test.ts b/src/test/terminals/serviceRegistry.unit.test.ts
@@ -14,17 +14,17 @@ import { TerminalCodeExecutionProvider } from '../../client/terminals/codeExecut
 import { TerminalDeactivateService } from '../../client/terminals/envCollectionActivation/deactivateService';
 import { TerminalIndicatorPrompt } from '../../client/terminals/envCollectionActivation/indicatorPrompt';
 import { TerminalEnvVarCollectionService } from '../../client/terminals/envCollectionActivation/service';
-import { ShellIntegrationService } from '../../client/terminals/envCollectionActivation/shellIntegrationService';
 import { registerTypes } from '../../client/terminals/serviceRegistry';
 import {
     ICodeExecutionHelper,
     ICodeExecutionManager,
     ICodeExecutionService,
-    IShellIntegrationService,
+    IShellIntegrationDetectionService,
     ITerminalAutoActivation,
     ITerminalDeactivateService,
     ITerminalEnvVarCollectionService,
 } from '../../client/terminals/types';
+import { ShellIntegrationDetectionService } from '../../client/terminals/envCollectionActivation/shellIntegrationService';
 
 suite('Terminal - Service Registry', () => {
     test('Ensure all services get registered', () => {
@@ -39,7 +39,7 @@ suite('Terminal - Service Registry', () => {
             [ITerminalEnvVarCollectionService, TerminalEnvVarCollectionService],
             [IExtensionSingleActivationService, TerminalIndicatorPrompt],
             [ITerminalDeactivateService, TerminalDeactivateService],
-            [IShellIntegrationService, ShellIntegrationService],
+            [IShellIntegrationDetectionService, ShellIntegrationDetectionService],
         ].forEach((args) => {
             if (args.length === 2) {
                 services