Auto select pipenv and poetry environments created for a workspace (#23102)

For #19153 

Filter out any pipenv or poetry envs which do not belong to the current
workspace.

Author: Kartik Raj

src/client/interpreter/configuration/environmentTypeComparer.ts

@@ -6,7 +6,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';
@@ -131,8 +136,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) {

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

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)) {

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;

src/client/pythonEnvironments/base/locator.ts

@@ -144,6 +144,7 @@ export type BasicEnvInfo = {
     executablePath: string;
     source?: PythonEnvSource[];
     envPath?: string;
+    searchLocation?: Uri;
 };
 
 /**

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.

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) {

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

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

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