Auto select pipenv and poetry environments created for a workspace (microsoft/vscode-python#23102)

Kartik Raj · web-flow · commit 89ac7ed0c6b0 · 2024-03-20T19:15:43.000-05:00
For microsoft/vscode-python#19153 Filter out any pipenv or poetry envs which do not belong to the current workspace.
diff --git a/extensions/positron-python/src/client/interpreter/configuration/environmentTypeComparer.ts b/extensions/positron-python/src/client/interpreter/configuration/environmentTypeComparer.ts
@@ -11,7 +11,12 @@ import { Resource } from '../../common/types';
 import { Architecture } from '../../common/utils/platform';
 import { isActiveStateEnvironmentForWorkspace } from '../../pythonEnvironments/common/environmentManagers/activestate';
 import { isParentPath } from '../../pythonEnvironments/common/externalDependencies';
-import { EnvironmentType, PythonEnvironment, virtualEnvTypes } from '../../pythonEnvironments/info';
+import {
+    EnvironmentType,
+    PythonEnvironment,
+    virtualEnvTypes,
+    workspaceVirtualEnvTypes,
+} from '../../pythonEnvironments/info';
 import { PythonVersion } from '../../pythonEnvironments/info/pythonVersion';
 import { IInterpreterHelper } from '../contracts';
 import { IInterpreterComparer } from './types';
@@ -147,8 +152,8 @@ export class EnvironmentTypeComparer implements IInterpreterComparer {
             if (getEnvLocationHeuristic(i, workspaceUri?.folderUri.fsPath || '') === EnvLocationHeuristic.Local) {
                 return true;
             }
-            if (virtualEnvTypes.includes(i.envType)) {
-                // We're not sure if these envs were created for the workspace, so do not recommend them.
+            if (!workspaceVirtualEnvTypes.includes(i.envType) && virtualEnvTypes.includes(i.envType)) {
+                // These are global virtual envs so we're not sure if these envs were created for the workspace, skip them.
                 return false;
             }
             if (i.version?.major === 2) {
diff --git a/extensions/positron-python/src/client/interpreter/configuration/interpreterSelector/interpreterSelector.ts b/extensions/positron-python/src/client/interpreter/configuration/interpreterSelector/interpreterSelector.ts
@@ -5,7 +5,7 @@
 
 import { inject, injectable } from 'inversify';
 import { Disposable, Uri } from 'vscode';
-import { arePathsSame } from '../../../common/platform/fs-paths';
+import { arePathsSame, isParentPath } from '../../../common/platform/fs-paths';
 import { IPathUtils, Resource } from '../../../common/types';
 import { getEnvPath } from '../../../pythonEnvironments/base/info/env';
 import { PythonEnvironment } from '../../../pythonEnvironments/info';
@@ -45,6 +45,13 @@ export class InterpreterSelector implements IInterpreterSelector {
         workspaceUri?: Uri,
         useDetailedName = false,
     ): IInterpreterQuickPickItem {
+        if (!useDetailedName) {
+            const workspacePath = workspaceUri?.fsPath;
+            if (workspacePath && isParentPath(interpreter.path, workspacePath)) {
+                // If interpreter is in the workspace, then display the full path.
+                useDetailedName = true;
+            }
+        }
         const path =
             interpreter.envPath && getEnvPath(interpreter.path, interpreter.envPath).pathType === 'envFolderPath'
                 ? interpreter.envPath
diff --git a/extensions/positron-python/src/client/pythonEnvironments/base/info/env.ts b/extensions/positron-python/src/client/pythonEnvironments/base/info/env.ts
@@ -160,7 +160,7 @@ export function setEnvDisplayString(env: PythonEnvInfo): void {
 
 function buildEnvDisplayString(env: PythonEnvInfo, getAllDetails = false): string {
     // main parts
-    const shouldDisplayKind = getAllDetails || env.searchLocation || globallyInstalledEnvKinds.includes(env.kind);
+    const shouldDisplayKind = getAllDetails || globallyInstalledEnvKinds.includes(env.kind);
     const shouldDisplayArch = !virtualEnvKinds.includes(env.kind);
     const displayNameParts: string[] = ['Python'];
     if (env.version && !isVersionEmpty(env.version)) {
diff --git a/extensions/positron-python/src/client/pythonEnvironments/base/info/index.ts b/extensions/positron-python/src/client/pythonEnvironments/base/info/index.ts
@@ -196,7 +196,7 @@ type _PythonEnvInfo = PythonEnvBaseInfo & PythonBuildInfo;
  * @prop distro - the installed Python distro that this env is using or belongs to
  * @prop display - the text to use when showing the env to users
  * @prop detailedDisplayName - display name containing all details
- * @prop searchLocation - the root under which a locator found this env, if any
+ * @prop searchLocation - the project to which this env is related to, if any
  */
 export type PythonEnvInfo = _PythonEnvInfo & {
     distro: PythonDistroInfo;
diff --git a/extensions/positron-python/src/client/pythonEnvironments/base/locator.ts b/extensions/positron-python/src/client/pythonEnvironments/base/locator.ts
@@ -144,6 +144,7 @@ export type BasicEnvInfo = {
     executablePath: string;
     source?: PythonEnvSource[];
     envPath?: string;
+    searchLocation?: Uri;
 };
 
 /**
diff --git a/extensions/positron-python/src/client/pythonEnvironments/base/locators/composite/envsReducer.ts b/extensions/positron-python/src/client/pythonEnvironments/base/locators/composite/envsReducer.ts
@@ -2,8 +2,9 @@
 // Licensed under the MIT License.
 
 import { cloneDeep, isEqual, uniq } from 'lodash';
-import { Event, EventEmitter } from 'vscode';
+import { Event, EventEmitter, Uri } from 'vscode';
 import { traceVerbose } from '../../../../logging';
+import { isParentPath } from '../../../common/externalDependencies';
 import { PythonEnvKind } from '../../info';
 import { areSameEnv } from '../../info/env';
 import { getPrioritizedEnvKinds } from '../../info/envKind';
@@ -136,9 +137,24 @@ function resolveEnvCollision(oldEnv: BasicEnvInfo, newEnv: BasicEnvInfo): BasicE
     const [env] = sortEnvInfoByPriority(oldEnv, newEnv);
     const merged = cloneDeep(env);
     merged.source = uniq((oldEnv.source ?? []).concat(newEnv.source ?? []));
+    merged.searchLocation = getMergedSearchLocation(oldEnv, newEnv);
     return merged;
 }
 
+function getMergedSearchLocation(oldEnv: BasicEnvInfo, newEnv: BasicEnvInfo): Uri | undefined {
+    if (oldEnv.searchLocation && newEnv.searchLocation) {
+        // Choose the deeper project path of the two, as that can be used to signify
+        // that the environment is related to both the projects.
+        if (isParentPath(oldEnv.searchLocation.fsPath, newEnv.searchLocation.fsPath)) {
+            return oldEnv.searchLocation;
+        }
+        if (isParentPath(newEnv.searchLocation.fsPath, oldEnv.searchLocation.fsPath)) {
+            return newEnv.searchLocation;
+        }
+    }
+    return oldEnv.searchLocation ?? newEnv.searchLocation;
+}
+
 /**
  * Selects an environment based on the environment selection priority. This should
  * match the priority in the environment identifier.
diff --git a/extensions/positron-python/src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts b/extensions/positron-python/src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts
@@ -53,11 +53,11 @@ function getResolvers(): Map<PythonEnvKind, (env: BasicEnvInfo) => Promise<Pytho
  * returned could still be invalid.
  */
 export async function resolveBasicEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
-    const { kind, source } = env;
+    const { kind, source, searchLocation } = env;
     const resolvers = getResolvers();
     const resolverForKind = resolvers.get(kind)!;
     const resolvedEnv = await resolverForKind(env);
-    resolvedEnv.searchLocation = getSearchLocation(resolvedEnv);
+    resolvedEnv.searchLocation = getSearchLocation(resolvedEnv, searchLocation);
     resolvedEnv.source = uniq(resolvedEnv.source.concat(source ?? []));
     if (getOSType() === OSType.Windows && resolvedEnv.source?.includes(PythonEnvSource.WindowsRegistry)) {
         // We can update env further using information we can get from the Windows registry.
@@ -87,7 +87,11 @@ async function getEnvType(env: PythonEnvInfo) {
     return undefined;
 }
 
-function getSearchLocation(env: PythonEnvInfo): Uri | undefined {
+function getSearchLocation(env: PythonEnvInfo, searchLocation: Uri | undefined): Uri | undefined {
+    if (searchLocation) {
+        // A search location has already been established by the downstream locators, simply use that.
+        return searchLocation;
+    }
     const folders = getWorkspaceFolderPaths();
     const isRootedEnv = folders.some((f) => isParentPath(env.executable.filename, f) || isParentPath(env.location, f));
     if (isRootedEnv) {
diff --git a/extensions/positron-python/src/client/pythonEnvironments/base/locators/lowLevel/globalVirtualEnvronmentLocator.ts b/extensions/positron-python/src/client/pythonEnvironments/base/locators/lowLevel/globalVirtualEnvronmentLocator.ts
@@ -3,14 +3,15 @@
 
 import { toLower, uniq, uniqBy } from 'lodash';
 import * as path from 'path';
+import { Uri } from 'vscode';
 import { chain, iterable } from '../../../../common/utils/async';
 import { getEnvironmentVariable, getOSType, getUserHomeDir, OSType } from '../../../../common/utils/platform';
 import { PythonEnvKind } from '../../info';
 import { BasicEnvInfo, IPythonEnvsIterator } from '../../locator';
 import { FSWatchingLocator } from './fsWatchingLocator';
 import { findInterpretersInDir, looksLikeBasicVirtualPython } from '../../../common/commonUtils';
 import { pathExists, untildify } from '../../../common/externalDependencies';
-import { isPipenvEnvironment } from '../../../common/environmentManagers/pipenv';
+import { getProjectDir, isPipenvEnvironment } from '../../../common/environmentManagers/pipenv';
 import {
     isVenvEnvironment,
     isVirtualenvEnvironment,
@@ -57,6 +58,18 @@ async function getGlobalVirtualEnvDirs(): Promise<string[]> {
     return [OSType.Windows, OSType.OSX].includes(getOSType()) ? uniqBy(venvDirs, toLower) : uniq(venvDirs);
 }
 
+async function getSearchLocation(env: BasicEnvInfo): Promise<Uri | undefined> {
+    if (env.kind === PythonEnvKind.Pipenv) {
+        // Pipenv environments are created only for a specific project, so they must only
+        // appear if that particular project is being queried.
+        const project = await getProjectDir(path.dirname(path.dirname(env.executablePath)));
+        if (project) {
+            return Uri.file(project);
+        }
+    }
+    return undefined;
+}
+
 /**
  * Gets the virtual environment kind for a given interpreter path.
  * This only checks for environments created using venv, virtualenv,
@@ -123,8 +136,9 @@ export class GlobalVirtualEnvironmentLocator extends FSWatchingLocator {
                             // check multiple times. Those checks are file system heavy and
                             // we can use the kind to determine this anyway.
                             const kind = await getVirtualEnvKind(filename);
+                            const searchLocation = await getSearchLocation({ kind, executablePath: filename });
                             try {
-                                yield { kind, executablePath: filename };
+                                yield { kind, executablePath: filename, searchLocation };
                                 traceVerbose(`Global Virtual Environment: [added] ${filename}`);
                             } catch (ex) {
                                 traceError(`Failed to process environment: ${filename}`, ex);
diff --git a/extensions/positron-python/src/client/pythonEnvironments/base/locators/lowLevel/poetryLocator.ts b/extensions/positron-python/src/client/pythonEnvironments/base/locators/lowLevel/poetryLocator.ts
@@ -4,6 +4,7 @@
 'use strict';
 
 import * as path from 'path';
+import { Uri } from 'vscode';
 import { chain, iterable } from '../../../../common/utils/async';
 import { PythonEnvKind } from '../../info';
 import { BasicEnvInfo, IPythonEnvsIterator } from '../../locator';
@@ -59,7 +60,7 @@ export class PoetryLocator extends LazyResourceBasedLocator {
                             // We should extract the kind here to avoid doing is*Environment()
                             // check multiple times. Those checks are file system heavy and
                             // we can use the kind to determine this anyway.
-                            yield { executablePath: filename, kind };
+                            yield { executablePath: filename, kind, searchLocation: Uri.file(root) };
                             traceVerbose(`Poetry Virtual Environment: [added] ${filename}`);
                         } catch (ex) {
                             traceError(`Failed to process environment: ${filename}`, ex);
diff --git a/extensions/positron-python/src/client/pythonEnvironments/common/environmentManagers/pipenv.ts b/extensions/positron-python/src/client/pythonEnvironments/common/environmentManagers/pipenv.ts
@@ -70,7 +70,7 @@ async function getPipfileIfLocal(interpreterPath: string): Promise<string | unde
  * Returns the project directory for pipenv environments given the environment folder
  * @param envFolder Path to the environment folder
  */
-async function getProjectDir(envFolder: string): Promise<string | undefined> {
+export async function getProjectDir(envFolder: string): Promise<string | undefined> {
     // Global pipenv environments have a .project file with the absolute path to the project
     // See https://github.com/pypa/pipenv/blob/v2018.6.25/CHANGELOG.rst#features--improvements
     // This is the layout we expect
diff --git a/extensions/positron-python/src/client/pythonEnvironments/info/index.ts b/extensions/positron-python/src/client/pythonEnvironments/info/index.ts
@@ -25,11 +25,14 @@ export enum EnvironmentType {
     Global = 'Global',
     System = 'System',
 }
+/**
+ * These envs are only created for a specific workspace, which we're able to detect.
+ */
+export const workspaceVirtualEnvTypes = [EnvironmentType.Poetry, EnvironmentType.Pipenv];
 
 export const virtualEnvTypes = [
-    EnvironmentType.Poetry,
-    EnvironmentType.Pipenv,
-    EnvironmentType.Hatch,
+    ...workspaceVirtualEnvTypes,
+    EnvironmentType.Hatch, // This is also a workspace virtual env, but we're not treating it as such as of today.
     EnvironmentType.Venv,
     EnvironmentType.VirtualEnvWrapper,
     EnvironmentType.Conda,
diff --git a/extensions/positron-python/src/test/pythonEnvironments/base/locators/composite/envsResolver.unit.test.ts b/extensions/positron-python/src/test/pythonEnvironments/base/locators/composite/envsResolver.unit.test.ts
@@ -57,7 +57,11 @@ suite('Python envs locator - Environments Resolver', () => {
     /**
      * Returns the expected environment to be returned by Environment info service
      */
-    function createExpectedEnvInfo(env: PythonEnvInfo, expectedDisplay: string): PythonEnvInfo {
+    function createExpectedEnvInfo(
+        env: PythonEnvInfo,
+        expectedDisplay: string,
+        expectedDetailedDisplay: string,
+    ): PythonEnvInfo {
         const updatedEnv = cloneDeep(env);
         updatedEnv.version = {
             ...parseVersion('3.8.3-final'),
@@ -67,7 +71,7 @@ suite('Python envs locator - Environments Resolver', () => {
         updatedEnv.executable.sysPrefix = 'path';
         updatedEnv.arch = Architecture.x64;
         updatedEnv.display = expectedDisplay;
-        updatedEnv.detailedDisplayName = expectedDisplay;
+        updatedEnv.detailedDisplayName = expectedDetailedDisplay;
         if (env.kind === PythonEnvKind.Conda) {
             env.type = PythonEnvType.Conda;
         }
@@ -82,6 +86,7 @@ suite('Python envs locator - Environments Resolver', () => {
         location = '',
         display: string | undefined = undefined,
         type?: PythonEnvType,
+        detailedDisplay?: string,
     ): PythonEnvInfo {
         return {
             name,
@@ -94,7 +99,7 @@ suite('Python envs locator - Environments Resolver', () => {
                 mtime: -1,
             },
             display,
-            detailedDisplayName: display,
+            detailedDisplayName: detailedDisplay ?? display,
             version,
             arch: Architecture.Unknown,
             distro: { org: '' },
@@ -134,8 +139,9 @@ suite('Python envs locator - Environments Resolver', () => {
                 undefined,
                 'win1',
                 path.join(testVirtualHomeDir, '.venvs', 'win1'),
-                "Python ('win1': venv)",
+                "Python ('win1')",
                 PythonEnvType.Virtual,
+                "Python ('win1': venv)",
             );
             const envsReturnedByParentLocator = [env1];
             const parentLocator = new SimpleLocator<BasicEnvInfo>(envsReturnedByParentLocator);
@@ -170,7 +176,11 @@ suite('Python envs locator - Environments Resolver', () => {
             const envs = await getEnvsWithUpdates(iterator);
 
             assertEnvsEqual(envs, [
-                createExpectedEnvInfo(resolvedEnvReturnedByBasicResolver, "Python 3.8.3 ('win1': venv)"),
+                createExpectedEnvInfo(
+                    resolvedEnvReturnedByBasicResolver,
+                    "Python 3.8.3 ('win1')",
+                    "Python 3.8.3 ('win1': venv)",
+                ),
             ]);
         });
 
@@ -237,7 +247,11 @@ suite('Python envs locator - Environments Resolver', () => {
 
             // Assert
             assertEnvsEqual(envs, [
-                createExpectedEnvInfo(resolvedUpdatedEnvReturnedByBasicResolver, "Python 3.8.3 ('win1': venv)"),
+                createExpectedEnvInfo(
+                    resolvedUpdatedEnvReturnedByBasicResolver,
+                    "Python 3.8.3 ('win1')",
+                    "Python 3.8.3 ('win1': venv)",
+                ),
             ]);
             didUpdate.dispose();
         });
@@ -377,7 +391,11 @@ suite('Python envs locator - Environments Resolver', () => {
 
             assertEnvEqual(
                 expected,
-                createExpectedEnvInfo(resolvedEnvReturnedByBasicResolver, "Python 3.8.3 ('win1': venv)"),
+                createExpectedEnvInfo(
+                    resolvedEnvReturnedByBasicResolver,
+                    "Python 3.8.3 ('win1')",
+                    "Python 3.8.3 ('win1': venv)",
+                ),
             );
         });
 
diff --git a/extensions/positron-python/src/test/pythonEnvironments/base/locators/envTestUtils.ts b/extensions/positron-python/src/test/pythonEnvironments/base/locators/envTestUtils.ts
@@ -103,10 +103,12 @@ export function assertBasicEnvsEqual(actualEnvs: BasicEnvInfo[], expectedEnvs: B
         const [actual, expected] = value;
         if (actual) {
             actual.source = actual.source ?? [];
+            actual.searchLocation = actual.searchLocation ?? undefined;
             actual.source.sort();
         }
         if (expected) {
             expected.source = expected.source ?? [];
+            expected.searchLocation = expected.searchLocation ?? undefined;
             expected.source.sort();
         }
         assert.deepStrictEqual(actual, expected);
diff --git a/extensions/positron-python/src/test/pythonEnvironments/base/locators/lowLevel/globalVirtualEnvironmentLocator.unit.test.ts b/extensions/positron-python/src/test/pythonEnvironments/base/locators/lowLevel/globalVirtualEnvironmentLocator.unit.test.ts
diff --git a/extensions/positron-python/src/test/pythonEnvironments/base/locators/lowLevel/poetryLocator.unit.test.ts b/extensions/positron-python/src/test/pythonEnvironments/base/locators/lowLevel/poetryLocator.unit.test.ts
