From 1f96862cfbed5b1f0c86ad0805deb3ef6a7b9727 Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 30 Aug 2023 19:01:45 +0000 Subject: [PATCH 1/3] Set PS1 for conda environments in non-Windows --- .../terminalEnvVarCollectionService.ts | 57 ++++++++++++++++--- 1 file changed, 48 insertions(+), 9 deletions(-) diff --git a/src/client/interpreter/activation/terminalEnvVarCollectionService.ts b/src/client/interpreter/activation/terminalEnvVarCollectionService.ts index 33d51e7e2ed0..3849160b7e4b 100644 --- a/src/client/interpreter/activation/terminalEnvVarCollectionService.ts +++ b/src/client/interpreter/activation/terminalEnvVarCollectionService.ts @@ -31,7 +31,7 @@ import { traceDecoratorVerbose, traceError, traceVerbose, traceWarn } from '../. import { IInterpreterService } from '../contracts'; import { defaultShells } from './service'; import { IEnvironmentActivationService, ITerminalEnvVarCollectionService } from './types'; -import { EnvironmentType } from '../../pythonEnvironments/info'; +import { EnvironmentType, PythonEnvironment } from '../../pythonEnvironments/info'; import { getSearchPathEnvVarNames } from '../../common/utils/exec'; import { EnvironmentVariables } from '../../common/variables/types'; import { TerminalShellType } from '../../common/terminal/types'; @@ -44,6 +44,15 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ virtualWorkspace: false, }; + /** + * Prompts for these shells cannot be set reliably using variables + */ + private noPromptVariableShells = [ + TerminalShellType.powershell, + TerminalShellType.powershellCore, + TerminalShellType.fish, + ]; + private deferred: Deferred | undefined; private registeredOnce = false; @@ -150,6 +159,10 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ ); } const processEnv = this.processEnvVars; + + // PS1 in some cases is a shell variable (not an env variable) so "env" might not contain it, calculate it in that case. + env.PS1 = await this.getPS1(shell, resource, env); + Object.keys(env).forEach((key) => { if (shouldSkip(key)) { return; @@ -213,15 +226,8 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ this.terminalPromptIsCorrect(resource); return; } - // Prompts for these shells cannot be set reliably using variables - const exceptionShells = [ - TerminalShellType.powershell, - TerminalShellType.powershellCore, - TerminalShellType.fish, - TerminalShellType.zsh, // TODO: Remove this once https://github.com/microsoft/vscode/issues/188875 is fixed - ]; const customShellType = identifyShellFromShellPath(shell); - if (exceptionShells.includes(customShellType)) { + if (this.noPromptVariableShells.includes(customShellType)) { return; } if (this.platform.osType !== OSType.Windows) { @@ -243,6 +249,26 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ this.terminalPromptIsCorrect(resource); } + private async getPS1(shell: string, resource: Resource, env: EnvironmentVariables) { + if (env.PS1) { + return env.PS1; + } + const customShellType = identifyShellFromShellPath(shell); + if (this.noPromptVariableShells.includes(customShellType)) { + return undefined; + } + if (this.platform.osType !== OSType.Windows) { + // These shells are expected to set PS1 variable for terminal prompt for virtual/conda environments. + const interpreter = await this.interpreterService.getActiveInterpreter(resource); + const shouldPS1BeSet = interpreter?.type !== undefined; + if (shouldPS1BeSet && !env.PS1) { + // PS1 should be set but no PS1 was set. + return getPromptForEnv(interpreter); + } + } + return undefined; + } + private async handleMicroVenv(resource: Resource) { try { const workspaceFolder = this.getWorkspaceFolder(resource); @@ -313,3 +339,16 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ function shouldSkip(env: string) { return ['_', 'SHLVL'].includes(env); } + +function getPromptForEnv(interpreter: PythonEnvironment | undefined) { + if (!interpreter) { + return undefined; + } + if (interpreter.envName) { + return `(${interpreter.envName}) `; + } + if (interpreter.envPath) { + return `(${path.basename(interpreter.envPath)}) `; + } + return undefined; +} From 548952d9b4da3087f74a1cba22a95779be22a483 Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 30 Aug 2023 19:08:19 +0000 Subject: [PATCH 2/3] Remove zsh safeguard test --- ...rminalEnvVarCollectionService.unit.test.ts | 26 ------------------- 1 file changed, 26 deletions(-) diff --git a/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts b/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts index 7393f4ad07ad..bf535513fa2d 100644 --- a/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts +++ b/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts @@ -321,32 +321,6 @@ suite('Terminal Environment Variable Collection Service', () => { expect(result).to.equal(false); }); - test('Correct track that prompt was not set for non-Windows zsh where PS1 is set', async () => { - when(platform.osType).thenReturn(OSType.Linux); - const envVars: NodeJS.ProcessEnv = { VIRTUAL_ENV: 'prefix/to/venv', PS1: '(.venv)', ...process.env }; - const ps1Shell = 'zsh'; - const resource = Uri.file('a'); - const workspaceFolder: WorkspaceFolder = { - uri: Uri.file('workspacePath'), - name: 'workspace1', - index: 0, - }; - when(interpreterService.getActiveInterpreter(resource)).thenResolve(({ - type: PythonEnvType.Virtual, - } as unknown) as PythonEnvironment); - when(workspaceService.getWorkspaceFolder(resource)).thenReturn(workspaceFolder); - when( - environmentActivationService.getActivatedEnvironmentVariables(resource, undefined, undefined, ps1Shell), - ).thenResolve(envVars); - when(collection.replace(anything(), anything(), anything())).thenReturn(); - - await terminalEnvVarCollectionService._applyCollection(resource, ps1Shell); - - const result = terminalEnvVarCollectionService.isTerminalPromptSetCorrectly(resource); - - expect(result).to.equal(false); - }); - test('Correct track that prompt was not set for non-Windows where PS1 is not set', async () => { when(platform.osType).thenReturn(OSType.Linux); const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env }; From 02760b59ccec5a868c4216e2cb005bb0d6e752f6 Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 30 Aug 2023 20:42:58 +0000 Subject: [PATCH 3/3] Add tests --- ...rminalEnvVarCollectionService.unit.test.ts | 32 ++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts b/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts index bf535513fa2d..2e2327bd181c 100644 --- a/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts +++ b/src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts @@ -321,7 +321,7 @@ suite('Terminal Environment Variable Collection Service', () => { expect(result).to.equal(false); }); - test('Correct track that prompt was not set for non-Windows where PS1 is not set', async () => { + test('Correct track that prompt was set for non-Windows where PS1 is not set but should be set', async () => { when(platform.osType).thenReturn(OSType.Linux); const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env }; const ps1Shell = 'zsh'; @@ -333,6 +333,36 @@ suite('Terminal Environment Variable Collection Service', () => { }; when(interpreterService.getActiveInterpreter(resource)).thenResolve(({ type: PythonEnvType.Conda, + envName: 'envName', + envPath: 'prefix/to/conda', + } as unknown) as PythonEnvironment); + when(workspaceService.getWorkspaceFolder(resource)).thenReturn(workspaceFolder); + when( + environmentActivationService.getActivatedEnvironmentVariables(resource, undefined, undefined, ps1Shell), + ).thenResolve(envVars); + when(collection.replace(anything(), anything(), anything())).thenReturn(); + + await terminalEnvVarCollectionService._applyCollection(resource, ps1Shell); + + const result = terminalEnvVarCollectionService.isTerminalPromptSetCorrectly(resource); + + expect(result).to.equal(true); + }); + + test('Correct track that prompt was not set for non-Windows fish where PS1 is not set', async () => { + when(platform.osType).thenReturn(OSType.Linux); + const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env }; + const ps1Shell = 'fish'; + const resource = Uri.file('a'); + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('workspacePath'), + name: 'workspace1', + index: 0, + }; + when(interpreterService.getActiveInterpreter(resource)).thenResolve(({ + type: PythonEnvType.Conda, + envName: 'envName', + envPath: 'prefix/to/conda', } as unknown) as PythonEnvironment); when(workspaceService.getWorkspaceFolder(resource)).thenReturn(workspaceFolder); when(