Set PS1 for conda environments in non-Windows when in pythonTerminalEnvVarActivation experiment

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<void> | 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;
+}

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

@@ -321,9 +321,9 @@ 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 () => {
+    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 = { VIRTUAL_ENV: 'prefix/to/venv', PS1: '(.venv)', ...process.env };
+        const envVars: NodeJS.ProcessEnv = { CONDA_PREFIX: 'prefix/to/conda', ...process.env };
         const ps1Shell = 'zsh';
         const resource = Uri.file('a');
         const workspaceFolder: WorkspaceFolder = {
@@ -332,7 +332,9 @@ suite('Terminal Environment Variable Collection Service', () => {
             index: 0,
         };
         when(interpreterService.getActiveInterpreter(resource)).thenResolve(({
-            type: PythonEnvType.Virtual,
+            type: PythonEnvType.Conda,
+            envName: 'envName',
+            envPath: 'prefix/to/conda',
         } as unknown) as PythonEnvironment);
         when(workspaceService.getWorkspaceFolder(resource)).thenReturn(workspaceFolder);
         when(
@@ -344,13 +346,13 @@ suite('Terminal Environment Variable Collection Service', () => {
 
         const result = terminalEnvVarCollectionService.isTerminalPromptSetCorrectly(resource);
 
-        expect(result).to.equal(false);
+        expect(result).to.equal(true);
     });
 
-    test('Correct track that prompt was not set for non-Windows where PS1 is not set', async () => {
+    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 = 'zsh';
+        const ps1Shell = 'fish';
         const resource = Uri.file('a');
         const workspaceFolder: WorkspaceFolder = {
             uri: Uri.file('workspacePath'),
@@ -359,6 +361,8 @@ 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(