Skip to content

Commit

Permalink
Set PS1 for conda environments in non-Windows when in `pythonTerminal…
Browse files Browse the repository at this point in the history
  • Loading branch information
Kartik Raj committed Aug 30, 2023
1 parent 7d25ceb commit 44f5bf7
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 15 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand All @@ -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);
Expand Down Expand Up @@ -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;
}
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand All @@ -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(
Expand All @@ -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'),
Expand All @@ -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(
Expand Down

0 comments on commit 44f5bf7

Please sign in to comment.