From 3fcb3fadfbf049a9541f3dfacff44f66b06c4b6f Mon Sep 17 00:00:00 2001 From: Anthony Kim <62267334+anthonykim1@users.noreply.github.com> Date: Thu, 19 Sep 2024 13:21:16 -0400 Subject: [PATCH] Adding PYTHONSTARTUP with shell integration to environment variable collection (#24111) Resolves: https://github.com/microsoft/vscode-python/issues/23930 - setting to opt out of PYTHONSTARTUP injection. --------- Co-authored-by: Courtney Webster <60238438+cwebster-99@users.noreply.github.com> --- package.json | 6 +++++ package.nls.json | 1 + src/client/common/types.ts | 1 + src/client/common/vscodeApis/workspaceApis.ts | 8 ++++++ src/client/extensionActivation.ts | 3 +++ .../envCollectionActivation/service.ts | 17 +++++++----- .../shellIntegrationService.ts | 4 +-- src/client/terminals/pythonStartup.ts | 27 +++++++++++++++++++ src/client/terminals/serviceRegistry.ts | 10 ++++--- src/client/terminals/types.ts | 9 +++++-- ...rminalEnvVarCollectionService.unit.test.ts | 6 ++--- .../terminals/codeExecution/helper.test.ts | 7 ++++- .../terminals/codeExecution/smartSend.test.ts | 7 ++++- .../terminals/serviceRegistry.unit.test.ts | 6 ++--- 14 files changed, 91 insertions(+), 21 deletions(-) create mode 100644 src/client/terminals/pythonStartup.ts diff --git a/package.json b/package.json index 696856aba561..af945dc2754e 100644 --- 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 index 5a5029231b17..f032f3d7c275 100644 --- 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 index 754e08004213..283319fd6cec 100644 --- 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 index 137382dc71a0..cb516da73075 100644 --- 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 { + return vscode.workspace.fs.createDirectory(uri); +} + +export function copy(source: vscode.Uri, dest: vscode.Uri, options?: { overwrite?: boolean }): Thenable { + return vscode.workspace.fs.copy(source, dest, options); +} diff --git a/src/client/extensionActivation.ts b/src/client/extensionActivation.ts index fe5d18a8b83f..429004e951cb 100644 --- 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).register(); + await registerPythonStartup(ext.context); + serviceManager.get(ICodeExecutionManager).registerCommands(); disposables.push(new ReplProvider(serviceContainer)); diff --git a/src/client/terminals/envCollectionActivation/service.ts b/src/client/terminals/envCollectionActivation/service.ts index 77e478b3577d..62971aa1fa98 100644 --- 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 { - 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 index 8ab3d84122b7..71cfb18dd437 100644 --- 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(); private readonly didChange = new EventEmitter(); diff --git a/src/client/terminals/pythonStartup.ts b/src/client/terminals/pythonStartup.ts new file mode 100644 index 000000000000..9a6b956d7f6e --- /dev/null +++ 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 { + const config = getConfiguration('python'); + const pythonrcSetting = config.get('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 index 3474edadd744..e62701dcec0e 100644 --- a/src/client/terminals/serviceRegistry.ts +++ b/src/client/terminals/serviceRegistry.ts @@ -12,7 +12,7 @@ import { ICodeExecutionHelper, ICodeExecutionManager, ICodeExecutionService, - IShellIntegrationService, + IShellIntegrationDetectionService, ITerminalAutoActivation, ITerminalDeactivateService, ITerminalEnvVarCollectionService, @@ -20,8 +20,8 @@ import { 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, CodeExecutionHelper); @@ -50,6 +50,10 @@ export function registerTypes(serviceManager: IServiceManager): void { IExtensionSingleActivationService, TerminalIndicatorPrompt, ); - serviceManager.addSingleton(IShellIntegrationService, ShellIntegrationService); + serviceManager.addSingleton( + IShellIntegrationDetectionService, + ShellIntegrationDetectionService, + ); + serviceManager.addBinding(ITerminalEnvVarCollectionService, IExtensionActivationService); } diff --git a/src/client/terminals/types.ts b/src/client/terminals/types.ts index 4c73da63dd1e..ada3acd851a9 100644 --- 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; isWorking(): Promise; } @@ -53,3 +53,8 @@ export interface ITerminalDeactivateService { initializeScriptParams(shell: string): Promise; getScriptLocation(shell: string, resource: Resource): Promise; } + +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 index 5d1027d12702..3550a92ba1ec 100644 --- 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(); shell = mock(); const envVarProvider = mock(); - shellIntegrationService = mock(); + shellIntegrationService = mock(); when(shellIntegrationService.isWorking()).thenResolve(true); globalCollection = mock(); collection = mock(); diff --git a/src/test/terminals/codeExecution/helper.test.ts b/src/test/terminals/codeExecution/helper.test.ts index 9a6deefcb7bf..e15c41957726 100644 --- 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 index 89f5ac2b5e4d..594db361f51e 100644 --- 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 index cb5b7715c4b7..4f865cdedc0d 100644 --- 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