Skip to content

Commit

Permalink
Select pyenv environment based on folder .python-version file (#23094)
Browse files Browse the repository at this point in the history
  • Loading branch information
Kartik Raj authored Mar 18, 2024
1 parent e9cd888 commit b9109cf
Show file tree
Hide file tree
Showing 5 changed files with 93 additions and 3 deletions.
1 change: 1 addition & 0 deletions src/client/interpreter/autoSelection/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,7 @@ export class InterpreterAutoSelectionService implements IInterpreterAutoSelectio
});
}

await this.envTypeComparer.initialize(resource);
const inExperiment = this.experimentService.inExperimentSync(DiscoveryUsingWorkers.experiment);
const workspaceUri = this.interpreterHelper.getActiveWorkspaceUri(resource);
let recommendedInterpreter: PythonEnvironment | undefined;
Expand Down
34 changes: 33 additions & 1 deletion src/client/interpreter/configuration/environmentTypeComparer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ import { EnvironmentType, PythonEnvironment, virtualEnvTypes } from '../../pytho
import { PythonVersion } from '../../pythonEnvironments/info/pythonVersion';
import { IInterpreterHelper } from '../contracts';
import { IInterpreterComparer } from './types';
import { getActivePyenvForDirectory } from '../../pythonEnvironments/common/environmentManagers/pyenv';
import { arePathsSame } from '../../common/platform/fs-paths';

export enum EnvLocationHeuristic {
/**
Expand All @@ -26,6 +28,8 @@ export enum EnvLocationHeuristic {
export class EnvironmentTypeComparer implements IInterpreterComparer {
private workspaceFolderPath: string;

private preferredPyenvInterpreterPath = new Map<string, string | undefined>();

constructor(@inject(IInterpreterHelper) private readonly interpreterHelper: IInterpreterHelper) {
this.workspaceFolderPath = this.interpreterHelper.getActiveWorkspaceUri(undefined)?.folderUri.fsPath ?? '';
}
Expand Down Expand Up @@ -54,6 +58,18 @@ export class EnvironmentTypeComparer implements IInterpreterComparer {
return envLocationComparison;
}

if (a.envType === EnvironmentType.Pyenv && b.envType === EnvironmentType.Pyenv) {
const preferredPyenv = this.preferredPyenvInterpreterPath.get(this.workspaceFolderPath);
if (preferredPyenv) {
if (arePathsSame(preferredPyenv, b.path)) {
return 1;
}
if (arePathsSame(preferredPyenv, a.path)) {
return -1;
}
}
}

// Check environment type.
const envTypeComparison = compareEnvironmentType(a, b);
if (envTypeComparison !== 0) {
Expand Down Expand Up @@ -85,6 +101,16 @@ export class EnvironmentTypeComparer implements IInterpreterComparer {
return nameA > nameB ? 1 : -1;
}

public async initialize(resource: Resource): Promise<void> {
const workspaceUri = this.interpreterHelper.getActiveWorkspaceUri(resource);
const cwd = workspaceUri?.folderUri.fsPath;
if (!cwd) {
return;
}
const preferredPyenvInterpreter = await getActivePyenvForDirectory(cwd);
this.preferredPyenvInterpreterPath.set(cwd, preferredPyenvInterpreter);
}

public getRecommended(interpreters: PythonEnvironment[], resource: Resource): PythonEnvironment | undefined {
// When recommending an intepreter for a workspace, we either want to return a local one
// or fallback on a globally-installed interpreter, and we don't want want to suggest a global environment
Expand Down Expand Up @@ -235,7 +261,13 @@ export function getEnvLocationHeuristic(environment: PythonEnvironment, workspac
*/
function compareEnvironmentType(a: PythonEnvironment, b: PythonEnvironment): number {
if (!a.type && !b.type) {
// Return 0 if two global interpreters are being compared.
// Unless one of them is pyenv interpreter, return 0 if two global interpreters are being compared.
if (a.envType === EnvironmentType.Pyenv && b.envType !== EnvironmentType.Pyenv) {
return -1;
}
if (a.envType !== EnvironmentType.Pyenv && b.envType === EnvironmentType.Pyenv) {
return 1;
}
return 0;
}
const envTypeByPriority = getPrioritizedEnvironmentType();
Expand Down
1 change: 1 addition & 0 deletions src/client/interpreter/configuration/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ export interface ISpecialQuickPickItem extends QuickPickItem {

export const IInterpreterComparer = Symbol('IInterpreterComparer');
export interface IInterpreterComparer {
initialize(resource: Resource): Promise<void>;
compare(a: PythonEnvironment, b: PythonEnvironment): number;
getRecommended(interpreters: PythonEnvironment[], resource: Resource): PythonEnvironment | undefined;
}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import * as path from 'path';
import { getEnvironmentVariable, getOSType, getUserHomeDir, OSType } from '../../../common/utils/platform';
import { arePathsSame, isParentPath, pathExists } from '../externalDependencies';
import { arePathsSame, isParentPath, pathExists, shellExecute } from '../externalDependencies';
import { traceVerbose } from '../../../logging';

export function getPyenvDir(): string {
// Check if the pyenv environment variables exist: PYENV on Windows, PYENV_ROOT on Unix.
Expand All @@ -20,6 +21,29 @@ export function getPyenvDir(): string {
return pyenvDir;
}

async function getPyenvBinary(): Promise<string | undefined> {
const pyenvDir = getPyenvDir();
const pyenvBin = path.join(pyenvDir, 'bin', 'pyenv');
if (await pathExists(pyenvBin)) {
return pyenvBin;
}
return undefined;
}

export async function getActivePyenvForDirectory(cwd: string): Promise<string | undefined> {
const pyenvBin = await getPyenvBinary();
if (!pyenvBin) {
return undefined;
}
try {
const pyenvInterpreterPath = await shellExecute(`${pyenvBin} which python`, { cwd });
return pyenvInterpreterPath.stdout.trim();
} catch (ex) {
traceVerbose(ex);
return undefined;
}
}

export function getPyenvVersionsDir(): string {
return path.join(getPyenvDir(), 'versions');
}
Expand Down
34 changes: 33 additions & 1 deletion src/test/configuration/environmentTypeComparer.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,15 @@ import {
} from '../../client/interpreter/configuration/environmentTypeComparer';
import { IInterpreterHelper } from '../../client/interpreter/contracts';
import { PythonEnvType } from '../../client/pythonEnvironments/base/info';
import * as pyenv from '../../client/pythonEnvironments/common/environmentManagers/pyenv';
import { EnvironmentType, PythonEnvironment } from '../../client/pythonEnvironments/info';

suite('Environment sorting', () => {
const workspacePath = path.join('path', 'to', 'workspace');
let interpreterHelper: IInterpreterHelper;
let getActiveWorkspaceUriStub: sinon.SinonStub;
let getInterpreterTypeDisplayNameStub: sinon.SinonStub;
const preferredPyenv = path.join('path', 'to', 'preferred', 'pyenv');

setup(() => {
getActiveWorkspaceUriStub = sinon.stub().returns({ folderUri: { fsPath: workspacePath } });
Expand All @@ -28,6 +30,8 @@ suite('Environment sorting', () => {
getActiveWorkspaceUri: getActiveWorkspaceUriStub,
getInterpreterTypeDisplayName: getInterpreterTypeDisplayNameStub,
} as unknown) as IInterpreterHelper;
const getActivePyenvForDirectory = sinon.stub(pyenv, 'getActivePyenvForDirectory');
getActivePyenvForDirectory.resolves(preferredPyenv);
});

teardown(() => {
Expand Down Expand Up @@ -147,6 +151,33 @@ suite('Environment sorting', () => {
} as PythonEnvironment,
expected: 1,
},
{
title: 'Preferred Pyenv interpreter should come before any global interpreter',
envA: {
envType: EnvironmentType.Pyenv,
version: { major: 3, minor: 12, patch: 2 },
path: preferredPyenv,
} as PythonEnvironment,
envB: {
envType: EnvironmentType.Pyenv,
version: { major: 3, minor: 10, patch: 2 },
path: path.join('path', 'to', 'normal', 'pyenv'),
} as PythonEnvironment,
expected: -1,
},
{
title: 'Pyenv interpreters should come first when there are global interpreters',
envA: {
envType: EnvironmentType.Global,
version: { major: 3, minor: 10, patch: 2 },
} as PythonEnvironment,
envB: {
envType: EnvironmentType.Pyenv,
version: { major: 3, minor: 7, patch: 2 },
path: path.join('path', 'to', 'normal', 'pyenv'),
} as PythonEnvironment,
expected: 1,
},
{
title: 'Global environment should not come first when there are global envs',
envA: {
Expand Down Expand Up @@ -283,8 +314,9 @@ suite('Environment sorting', () => {
];

testcases.forEach(({ title, envA, envB, expected }) => {
test(title, () => {
test(title, async () => {
const envTypeComparer = new EnvironmentTypeComparer(interpreterHelper);
await envTypeComparer.initialize(undefined);
const result = envTypeComparer.compare(envA, envB);

assert.strictEqual(result, expected);
Expand Down

0 comments on commit b9109cf

Please sign in to comment.