Select pyenv environment based on folder .python-version file (#23094)

Author: Kartik Raj

src/client/interpreter/autoSelection/index.ts

@@ -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;

src/client/interpreter/configuration/environmentTypeComparer.ts

@@ -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 {
     /**
@@ -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 ?? '';
     }
@@ -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) {
@@ -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
@@ -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();

src/client/interpreter/configuration/types.ts

@@ -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;
 }

src/client/pythonEnvironments/common/environmentManagers/pyenv.ts

@@ -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.
@@ -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');
 }

src/test/configuration/environmentTypeComparer.unit.test.ts

@@ -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 } });
@@ -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(() => {
@@ -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: {
@@ -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);