Skip to content

Commit 44f5bf7

Browse files
author
Kartik Raj
authored
Set PS1 for conda environments in non-Windows when in pythonTerminalEnvVarActivation experiment (#21902)
For #20822 ![image](https://github.com/microsoft/vscode-python/assets/13199757/8c9d4c87-54f2-4661-b6c6-c3b49ee3ff7a)
1 parent 7d25ceb commit 44f5bf7

File tree

2 files changed

+58
-15
lines changed

2 files changed

+58
-15
lines changed

src/client/interpreter/activation/terminalEnvVarCollectionService.ts

+48-9
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ import { traceDecoratorVerbose, traceError, traceVerbose, traceWarn } from '../.
3131
import { IInterpreterService } from '../contracts';
3232
import { defaultShells } from './service';
3333
import { IEnvironmentActivationService, ITerminalEnvVarCollectionService } from './types';
34-
import { EnvironmentType } from '../../pythonEnvironments/info';
34+
import { EnvironmentType, PythonEnvironment } from '../../pythonEnvironments/info';
3535
import { getSearchPathEnvVarNames } from '../../common/utils/exec';
3636
import { EnvironmentVariables } from '../../common/variables/types';
3737
import { TerminalShellType } from '../../common/terminal/types';
@@ -44,6 +44,15 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
4444
virtualWorkspace: false,
4545
};
4646

47+
/**
48+
* Prompts for these shells cannot be set reliably using variables
49+
*/
50+
private noPromptVariableShells = [
51+
TerminalShellType.powershell,
52+
TerminalShellType.powershellCore,
53+
TerminalShellType.fish,
54+
];
55+
4756
private deferred: Deferred<void> | undefined;
4857

4958
private registeredOnce = false;
@@ -150,6 +159,10 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
150159
);
151160
}
152161
const processEnv = this.processEnvVars;
162+
163+
// PS1 in some cases is a shell variable (not an env variable) so "env" might not contain it, calculate it in that case.
164+
env.PS1 = await this.getPS1(shell, resource, env);
165+
153166
Object.keys(env).forEach((key) => {
154167
if (shouldSkip(key)) {
155168
return;
@@ -213,15 +226,8 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
213226
this.terminalPromptIsCorrect(resource);
214227
return;
215228
}
216-
// Prompts for these shells cannot be set reliably using variables
217-
const exceptionShells = [
218-
TerminalShellType.powershell,
219-
TerminalShellType.powershellCore,
220-
TerminalShellType.fish,
221-
TerminalShellType.zsh, // TODO: Remove this once https://github.com/microsoft/vscode/issues/188875 is fixed
222-
];
223229
const customShellType = identifyShellFromShellPath(shell);
224-
if (exceptionShells.includes(customShellType)) {
230+
if (this.noPromptVariableShells.includes(customShellType)) {
225231
return;
226232
}
227233
if (this.platform.osType !== OSType.Windows) {
@@ -243,6 +249,26 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
243249
this.terminalPromptIsCorrect(resource);
244250
}
245251

252+
private async getPS1(shell: string, resource: Resource, env: EnvironmentVariables) {
253+
if (env.PS1) {
254+
return env.PS1;
255+
}
256+
const customShellType = identifyShellFromShellPath(shell);
257+
if (this.noPromptVariableShells.includes(customShellType)) {
258+
return undefined;
259+
}
260+
if (this.platform.osType !== OSType.Windows) {
261+
// These shells are expected to set PS1 variable for terminal prompt for virtual/conda environments.
262+
const interpreter = await this.interpreterService.getActiveInterpreter(resource);
263+
const shouldPS1BeSet = interpreter?.type !== undefined;
264+
if (shouldPS1BeSet && !env.PS1) {
265+
// PS1 should be set but no PS1 was set.
266+
return getPromptForEnv(interpreter);
267+
}
268+
}
269+
return undefined;
270+
}
271+
246272
private async handleMicroVenv(resource: Resource) {
247273
try {
248274
const workspaceFolder = this.getWorkspaceFolder(resource);
@@ -313,3 +339,16 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
313339
function shouldSkip(env: string) {
314340
return ['_', 'SHLVL'].includes(env);
315341
}
342+
343+
function getPromptForEnv(interpreter: PythonEnvironment | undefined) {
344+
if (!interpreter) {
345+
return undefined;
346+
}
347+
if (interpreter.envName) {
348+
return `(${interpreter.envName}) `;
349+
}
350+
if (interpreter.envPath) {
351+
return `(${path.basename(interpreter.envPath)}) `;
352+
}
353+
return undefined;
354+
}

src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts

+10-6
Original file line numberDiff line numberDiff line change
@@ -321,9 +321,9 @@ suite('Terminal Environment Variable Collection Service', () => {
321321
expect(result).to.equal(false);
322322
});
323323

324-
test('Correct track that prompt was not set for non-Windows zsh where PS1 is set', async () => {
324+
test('Correct track that prompt was set for non-Windows where PS1 is not set but should be set', async () => {
325325
when(platform.osType).thenReturn(OSType.Linux);
326-
const envVars: NodeJS.ProcessEnv = { VIRTUAL_ENV: 'prefix/to/venv', PS1: '(.venv)', ...process.env };
326+
const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env };
327327
const ps1Shell = 'zsh';
328328
const resource = Uri.file('a');
329329
const workspaceFolder: WorkspaceFolder = {
@@ -332,7 +332,9 @@ suite('Terminal Environment Variable Collection Service', () => {
332332
index: 0,
333333
};
334334
when(interpreterService.getActiveInterpreter(resource)).thenResolve(({
335-
type: PythonEnvType.Virtual,
335+
type: PythonEnvType.Conda,
336+
envName: 'envName',
337+
envPath: 'prefix/to/conda',
336338
} as unknown) as PythonEnvironment);
337339
when(workspaceService.getWorkspaceFolder(resource)).thenReturn(workspaceFolder);
338340
when(
@@ -344,13 +346,13 @@ suite('Terminal Environment Variable Collection Service', () => {
344346

345347
const result = terminalEnvVarCollectionService.isTerminalPromptSetCorrectly(resource);
346348

347-
expect(result).to.equal(false);
349+
expect(result).to.equal(true);
348350
});
349351

350-
test('Correct track that prompt was not set for non-Windows where PS1 is not set', async () => {
352+
test('Correct track that prompt was not set for non-Windows fish where PS1 is not set', async () => {
351353
when(platform.osType).thenReturn(OSType.Linux);
352354
const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env };
353-
const ps1Shell = 'zsh';
355+
const ps1Shell = 'fish';
354356
const resource = Uri.file('a');
355357
const workspaceFolder: WorkspaceFolder = {
356358
uri: Uri.file('workspacePath'),
@@ -359,6 +361,8 @@ suite('Terminal Environment Variable Collection Service', () => {
359361
};
360362
when(interpreterService.getActiveInterpreter(resource)).thenResolve(({
361363
type: PythonEnvType.Conda,
364+
envName: 'envName',
365+
envPath: 'prefix/to/conda',
362366
} as unknown) as PythonEnvironment);
363367
when(workspaceService.getWorkspaceFolder(resource)).thenReturn(workspaceFolder);
364368
when(

0 commit comments

Comments
 (0)