Only show envs related to current worksapce

Author: Kartik Raj

news/1 Enhancements/18263.md

@@ -0,0 +1 @@
+Only display Pipenv and Poetry envs related to the currently opened workspace in "Select Interpreter" UI.

src/client/pythonEnvironments/base/info/index.ts

@@ -168,7 +168,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 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

@@ -116,7 +116,12 @@ export type PythonLocatorQuery = BasicPythonLocatorQuery & {
 
 type QueryForEvent<E> = E extends PythonEnvsChangedEvent ? PythonLocatorQuery : BasicPythonLocatorQuery;
 
-export type BasicEnvInfo = { kind: PythonEnvKind; executablePath: string; source?: PythonEnvSource[] };
+export type BasicEnvInfo = {
+    kind: PythonEnvKind;
+    executablePath: string;
+    source?: PythonEnvSource[];
+    searchLocation?: Uri;
+};
 
 /**
  * A single Python environment locator.

src/client/pythonEnvironments/base/locators/composite/envsReducer.ts

@@ -119,6 +119,7 @@ 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 = oldEnv.searchLocation ?? newEnv.searchLocation;
     return merged;
 }
 

src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts

@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 
 import * as path from 'path';
-import { Uri } from 'vscode';
 import { uniq } from 'lodash';
 import { PythonEnvInfo, PythonEnvKind, PythonEnvSource, UNKNOWN_PYTHON_VERSION, virtualEnvKinds } from '../../info';
 import { buildEnvInfo, comparePythonVersionSpecificity, getEnvDisplayString, getEnvMatcher } from '../../info/env';
@@ -11,7 +10,7 @@ import {
     getInterpreterPathFromDir,
     getPythonVersionFromPath,
 } from '../../../common/commonUtils';
-import { arePathsSame, getWorkspaceFolders, isParentPath } from '../../../common/externalDependencies';
+import { arePathsSame } from '../../../common/externalDependencies';
 import { AnacondaCompanyName, Conda } from '../../../common/environmentManagers/conda';
 import { parsePyenvVersion } from '../../../common/environmentManagers/pyenv';
 import { Architecture, getOSType, OSType } from '../../../../common/utils/platform';
@@ -42,11 +41,16 @@ function getResolvers(): Map<PythonEnvKind, (executablePath: string) => Promise<
  * executable and returns it. Notice `undefined` is never returned, so environment
  * returned could still be invalid.
  */
-export async function resolveBasicEnv({ kind, executablePath, source }: BasicEnvInfo): Promise<PythonEnvInfo> {
+export async function resolveBasicEnv({
+    kind,
+    executablePath,
+    source,
+    searchLocation,
+}: BasicEnvInfo): Promise<PythonEnvInfo> {
     const resolvers = getResolvers();
     const resolverForKind = resolvers.get(kind)!;
     const resolvedEnv = await resolverForKind(executablePath);
-    resolvedEnv.searchLocation = getSearchLocation(resolvedEnv);
+    resolvedEnv.searchLocation = 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.
@@ -56,24 +60,6 @@ export async function resolveBasicEnv({ kind, executablePath, source }: BasicEnv
     return resolvedEnv;
 }
 
-function getSearchLocation(env: PythonEnvInfo): Uri | undefined {
-    const folders = getWorkspaceFolders();
-    const isRootedEnv = folders.some((f) => isParentPath(env.executable.filename, f));
-    if (isRootedEnv) {
-        // For environments inside roots, we need to set search location so they can be queried accordingly.
-        // Search location particularly for virtual environments is intended as the directory in which the
-        // environment was found in.
-        // For eg.the default search location for an env containing 'bin' or 'Scripts' directory is:
-        //
-        // searchLocation <--- Default search location directory
-        // |__ env
-        //    |__ bin or Scripts
-        //        |__ python  <--- executable
-        return Uri.file(path.dirname(env.location));
-    }
-    return undefined;
-}
-
 async function updateEnvUsingRegistry(env: PythonEnvInfo): Promise<void> {
     // Environment source has already been identified as windows registry, so we expect windows registry
     // cache to already be populated. Call sync function which relies on cache.

src/client/pythonEnvironments/base/locators/lowLevel/globalVirtualEnvronmentLocator.ts

@@ -3,14 +3,19 @@
 
 import { uniq } 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 {
+    findInterpretersInDir,
+    getEnvironmentDirFromPath,
+    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,
@@ -79,6 +84,18 @@ async function getVirtualEnvKind(interpreterPath: string): Promise<PythonEnvKind
     return PythonEnvKind.Unknown;
 }
 
+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(getEnvironmentDirFromPath(env.executablePath));
+        if (project) {
+            return Uri.file(project);
+        }
+    }
+    return undefined;
+}
+
 /**
  * Finds and resolves virtual environments created in known global locations.
  */
@@ -118,7 +135,8 @@ export class GlobalVirtualEnvironmentLocator extends FSWatchingLocator<BasicEnvI
                             // we can use the kind to determine this anyway.
                             const kind = await getVirtualEnvKind(filename);
                             try {
-                                yield { kind, executablePath: filename };
+                                const searchLocation = await getSearchLocation({ 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';
@@ -83,7 +84,7 @@ export class PoetryLocator extends FSWatchingLocator<BasicEnvInfo> {
                             // 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/base/locators/lowLevel/workspaceVirtualEnvLocator.ts

@@ -2,8 +2,13 @@
 // Licensed under the MIT License.
 
 import * as path from 'path';
+import { Uri } from 'vscode';
 import { chain, iterable } from '../../../../common/utils/async';
-import { findInterpretersInDir, looksLikeBasicVirtualPython } from '../../../common/commonUtils';
+import {
+    findInterpretersInDir,
+    getEnvironmentDirFromPath,
+    looksLikeBasicVirtualPython,
+} from '../../../common/commonUtils';
 import { pathExists } from '../../../common/externalDependencies';
 import { isPipenvEnvironment } from '../../../common/environmentManagers/pipenv';
 import { isVenvEnvironment, isVirtualenvEnvironment } from '../../../common/environmentManagers/simplevirtualenvs';
@@ -80,11 +85,19 @@ export class WorkspaceVirtualEnvironmentLocator extends FSWatchingLocator<BasicE
                         // python.exe or python in the same directory in the case of virtual
                         // environments.
                         if (await looksLikeBasicVirtualPython(entry)) {
-                            // 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.
+                            const location = getEnvironmentDirFromPath(filename);
+                            // For environments inside roots, we need to set search location so they can be queried accordingly.
+                            // Search location particularly for virtual environments is intended as the directory in which the
+                            // environment was found in.
+                            // For eg.the default search location for an env containing 'bin' or 'Scripts' directory is:
+                            //
+                            // searchLocation <--- Default search location directory
+                            // |__ env
+                            //    |__ bin or Scripts
+                            //        |__ python  <--- executable
+                            const searchLocation = Uri.file(path.dirname(location));
                             const kind = await getVirtualEnvKind(filename);
-                            yield { kind, executablePath: filename };
+                            yield { kind, executablePath: filename, searchLocation };
                             traceVerbose(`Workspace Virtual Environment: [added] ${filename}`);
                         } else {
                             traceVerbose(`Workspace Virtual Environment: [skipped] ${filename}`);

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

@@ -4,6 +4,7 @@
 import * as path from 'path';
 import { getEnvironmentVariable } from '../../../common/utils/platform';
 import { traceError } from '../../../logging';
+import { getEnvironmentDirFromPath } from '../commonUtils';
 import { arePathsSame, normCasePath, pathExists, readFile } from '../externalDependencies';
 
 function getSearchHeight() {
@@ -70,7 +71,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
@@ -98,7 +99,7 @@ async function getProjectDir(envFolder: string): Promise<string | undefined> {
  * @param interpreterPath Absolute path to any python interpreter.
  */
 async function getPipfileIfGlobal(interpreterPath: string): Promise<string | undefined> {
-    const envFolder = path.dirname(path.dirname(interpreterPath));
+    const envFolder = getEnvironmentDirFromPath(interpreterPath);
     const projectDir = await getProjectDir(envFolder);
     if (projectDir === undefined) {
         return undefined;