Ensure resolvers accept data from native locators (#23444)

DonJayamanne · DonJayamanne · commit 2e56a2150c20 · 2024-06-24T16:35:51.000+10:00
diff --git a/src/client/pythonEnvironments/base/info/env.ts b/src/client/pythonEnvironments/base/info/env.ts
@@ -42,6 +42,12 @@ export function buildEnvInfo(init?: {
     sysPrefix?: string;
     searchLocation?: Uri;
     type?: PythonEnvType;
+    /**
+     * Command used to run Python in this environment.
+     * E.g. `conda run -n envName python` or `python.exe`
+     */
+    pythonRunCommand?: string[];
+    identifiedUsingNativeLocator?: boolean;
 }): PythonEnvInfo {
     const env: PythonEnvInfo = {
         name: init?.name ?? '',
@@ -69,6 +75,8 @@ export function buildEnvInfo(init?: {
             org: init?.org ?? '',
         },
         source: init?.source ?? [],
+        pythonRunCommand: init?.pythonRunCommand,
+        identifiedUsingNativeLocator: init?.identifiedUsingNativeLocator,
     };
     if (init !== undefined) {
         updateEnv(env, init);
diff --git a/src/client/pythonEnvironments/base/info/index.ts b/src/client/pythonEnvironments/base/info/index.ts
@@ -203,6 +203,12 @@ export type PythonEnvInfo = _PythonEnvInfo & {
     display?: string;
     detailedDisplayName?: string;
     searchLocation?: Uri;
+    /**
+     * Command used to run Python in this environment.
+     * E.g. `conda run -n envName python` or `python.exe`
+     */
+    pythonRunCommand?: string[];
+    identifiedUsingNativeLocator?: boolean;
 };
 
 /**
diff --git a/src/client/pythonEnvironments/base/locator.ts b/src/client/pythonEnvironments/base/locator.ts
@@ -144,9 +144,24 @@ export type BasicEnvInfo = {
     executablePath: string;
     source?: PythonEnvSource[];
     envPath?: string;
+    /**
+     * The project to which this env is related to, if any
+     * E.g. the project directory when dealing with pipenv virtual environments.
+     */
     searchLocation?: Uri;
     version?: PythonVersion;
     name?: string;
+    /**
+     * Display name provided by locators, not generated by us.
+     * E.g. display name as provided by Windows Registry or Windows Store, etc
+     */
+    displayName?: string;
+    /**
+     * Command used to run Python in this environment.
+     * E.g. `conda run -n envName python` or `python.exe`
+     */
+    pythonRunCommand?: string[];
+    identifiedUsingNativeLocator?: boolean;
 };
 
 /**
diff --git a/src/client/pythonEnvironments/base/locators/common/nativePythonFinder.ts b/src/client/pythonEnvironments/base/locators/common/nativePythonFinder.ts
@@ -23,6 +23,7 @@ export interface NativeEnvInfo {
     version?: string;
     pythonRunCommand?: string[];
     envPath?: string;
+    envManager?: NativeEnvManagerInfo;
     /**
      * Path to the project directory when dealing with pipenv virtual environments.
      */
diff --git a/src/client/pythonEnvironments/base/locators/composite/envsResolver.ts b/src/client/pythonEnvironments/base/locators/composite/envsResolver.ts
@@ -154,6 +154,12 @@ export class PythonEnvsResolver implements IResolvingLocator {
 
 async function setKind(env: BasicEnvInfo, environmentKinds: Map<string, PythonEnvKind>) {
     const { path } = getEnvPath(env.executablePath, env.envPath);
+    // For native locators, do not try to identify the environment kind.
+    // its already set by the native locator & thats accurate.
+    if (env.identifiedUsingNativeLocator) {
+        environmentKinds.set(path, env.kind);
+        return;
+    }
     let kind = environmentKinds.get(path);
     if (!kind) {
         if (!isIdentifierRegistered(env.kind)) {
diff --git a/src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts b/src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts
@@ -139,14 +139,21 @@ async function resolveGloballyInstalledEnv(env: BasicEnvInfo): Promise<PythonEnv
     const { executablePath } = env;
     let version;
     try {
-        version = parseVersionFromExecutable(executablePath);
+        version = env.version ?? parseVersionFromExecutable(executablePath);
     } catch {
         version = UNKNOWN_PYTHON_VERSION;
     }
     const envInfo = buildEnvInfo({
         kind: env.kind,
+        name: env.name,
+        display: env.displayName,
+        sysPrefix: env.envPath,
+        location: env.envPath,
+        searchLocation: env.searchLocation,
         version,
         executable: executablePath,
+        pythonRunCommand: env.pythonRunCommand,
+        identifiedUsingNativeLocator: env.identifiedUsingNativeLocator,
     });
     return envInfo;
 }
@@ -155,17 +162,59 @@ async function resolveSimpleEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
     const { executablePath, kind } = env;
     const envInfo = buildEnvInfo({
         kind,
-        version: await getPythonVersionFromPath(executablePath),
+        version: env.version ?? (await getPythonVersionFromPath(executablePath)),
         executable: executablePath,
+        sysPrefix: env.envPath,
+        location: env.envPath,
+        display: env.displayName,
+        searchLocation: env.searchLocation,
+        pythonRunCommand: env.pythonRunCommand,
+        identifiedUsingNativeLocator: env.identifiedUsingNativeLocator,
+        name: env.name,
         type: PythonEnvType.Virtual,
     });
-    const location = getEnvironmentDirFromPath(executablePath);
+    const location = env.envPath ?? getEnvironmentDirFromPath(executablePath);
     envInfo.location = location;
     envInfo.name = path.basename(location);
     return envInfo;
 }
 
 async function resolveCondaEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
+    if (env.identifiedUsingNativeLocator) {
+        // New approach using native locator.
+        const executable = env.executablePath;
+        const envPath = env.envPath ?? getEnvironmentDirFromPath(executable);
+        // TODO: Hacky, `executable` is never undefined in the typedef,
+        // However, in reality with native locator this can be undefined.
+        const version = env.version ?? (executable ? await getPythonVersionFromPath(executable) : undefined);
+        const info = buildEnvInfo({
+            executable,
+            kind: PythonEnvKind.Conda,
+            org: AnacondaCompanyName,
+            location: envPath,
+            sysPrefix: envPath,
+            display: env.displayName,
+            pythonRunCommand: env.pythonRunCommand,
+            identifiedUsingNativeLocator: env.identifiedUsingNativeLocator,
+            searchLocation: env.searchLocation,
+            source: [],
+            version,
+            type: PythonEnvType.Conda,
+            name: env.name,
+        });
+
+        if (env.envPath && executable && path.basename(executable) === executable) {
+            // For environments without python, set ID using the predicted executable path after python is installed.
+            // Another alternative could've been to set ID of all conda environments to the environment path, as that
+            // remains constant even after python installation.
+            const predictedExecutable = getCondaInterpreterPath(env.envPath);
+            info.id = getEnvID(predictedExecutable, env.envPath);
+        }
+        return info;
+    }
+
+    // Old approach (without native locator).
+    // In this approach we need to find conda.
     const { executablePath } = env;
     const conda = await Conda.getConda();
     if (conda === undefined) {
@@ -210,18 +259,26 @@ async function resolveCondaEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
 
 async function resolvePyenvEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
     const { executablePath } = env;
-    const location = getEnvironmentDirFromPath(executablePath);
+    const location = env.envPath ?? getEnvironmentDirFromPath(executablePath);
     const name = path.basename(location);
 
     // The sub-directory name sometimes can contain distro and python versions.
     // here we attempt to extract the texts out of the name.
     const versionStrings = parsePyenvVersion(name);
 
     const envInfo = buildEnvInfo({
-        kind: PythonEnvKind.Pyenv,
+        // If using native resolver, then we can get the kind from the native resolver.
+        // E.g. pyenv can have conda environments as well.
+        kind: env.identifiedUsingNativeLocator && env.kind ? env.kind : PythonEnvKind.Pyenv,
         executable: executablePath,
         source: [],
         location,
+        searchLocation: env.searchLocation,
+        sysPrefix: env.envPath,
+        display: env.displayName,
+        name: env.name,
+        pythonRunCommand: env.pythonRunCommand,
+        identifiedUsingNativeLocator: env.identifiedUsingNativeLocator,
         // Pyenv environments can fall in to these three categories:
         // 1. Global Installs : These are environments that are created when you install
         //    a supported python distribution using `pyenv install <distro>` command.
@@ -240,14 +297,17 @@ async function resolvePyenvEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
         //
         // Here we look for near by files, or config files to see if we can get python version info
         // without running python itself.
-        version: await getPythonVersionFromPath(executablePath, versionStrings?.pythonVer),
+        version: env.version ?? (await getPythonVersionFromPath(executablePath, versionStrings?.pythonVer)),
         org: versionStrings && versionStrings.distro ? versionStrings.distro : '',
     });
 
-    if (await isBaseCondaPyenvEnvironment(executablePath)) {
-        envInfo.name = 'base';
-    } else {
-        envInfo.name = name;
+    // Do this only for the old approach, when not using native locators.
+    if (!env.identifiedUsingNativeLocator) {
+        if (await isBaseCondaPyenvEnvironment(executablePath)) {
+            envInfo.name = 'base';
+        } else {
+            envInfo.name = name;
+        }
     }
     return envInfo;
 }
@@ -256,6 +316,14 @@ async function resolveActiveStateEnv(env: BasicEnvInfo): Promise<PythonEnvInfo>
     const info = buildEnvInfo({
         kind: env.kind,
         executable: env.executablePath,
+        display: env.displayName,
+        version: env.version,
+        identifiedUsingNativeLocator: env.identifiedUsingNativeLocator,
+        location: env.envPath,
+        name: env.name,
+        pythonRunCommand: env.pythonRunCommand,
+        searchLocation: env.searchLocation,
+        sysPrefix: env.envPath,
     });
     const projects = await ActiveState.getState().then((v) => v?.getProjects());
     if (projects) {
@@ -285,8 +353,15 @@ async function resolveMicrosoftStoreEnv(env: BasicEnvInfo): Promise<PythonEnvInf
     return buildEnvInfo({
         kind: PythonEnvKind.MicrosoftStore,
         executable: executablePath,
-        version: parsePythonVersionFromPath(executablePath),
+        version: env.version ?? parsePythonVersionFromPath(executablePath),
         org: 'Microsoft',
+        display: env.displayName,
+        location: env.envPath,
+        sysPrefix: env.envPath,
+        searchLocation: env.searchLocation,
+        name: env.name,
+        pythonRunCommand: env.pythonRunCommand,
+        identifiedUsingNativeLocator: env.identifiedUsingNativeLocator,
         arch: Architecture.x64,
         source: [PythonEnvSource.PathEnvVar],
     });
diff --git a/src/client/pythonEnvironments/base/locators/lowLevel/nativeLocator.ts b/src/client/pythonEnvironments/base/locators/lowLevel/nativeLocator.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { Disposable, Event, EventEmitter } from 'vscode';
+import { Disposable, Event, EventEmitter, Uri } from 'vscode';
 import { IDisposable } from '../../../../common/types';
 import { ILocator, BasicEnvInfo, IPythonEnvsIterator } from '../../locator';
 import { PythonEnvsChangedEvent } from '../../watcher';
@@ -116,6 +116,10 @@ export class NativeLocator implements ILocator<BasicEnvInfo>, IDisposable {
                         envPath: data.envPath,
                         version: parseVersion(data.version),
                         name: data.name === '' ? undefined : data.name,
+                        displayName: data.displayName,
+                        pythonRunCommand: data.pythonRunCommand,
+                        searchLocation: data.projectPath ? Uri.file(data.projectPath) : undefined,
+                        identifiedUsingNativeLocator: true,
                     });
                 }
             }),
diff --git a/src/test/pythonEnvironments/base/locators/composite/envsResolver.unit.test.ts b/src/test/pythonEnvironments/base/locators/composite/envsResolver.unit.test.ts
@@ -72,6 +72,8 @@ suite('Python envs locator - Environments Resolver', () => {
         updatedEnv.arch = Architecture.x64;
         updatedEnv.display = expectedDisplay;
         updatedEnv.detailedDisplayName = expectedDetailedDisplay;
+        updatedEnv.identifiedUsingNativeLocator = updatedEnv.identifiedUsingNativeLocator ?? undefined;
+        updatedEnv.pythonRunCommand = updatedEnv.pythonRunCommand ?? undefined;
         if (env.kind === PythonEnvKind.Conda) {
             env.type = PythonEnvType.Conda;
         }
@@ -106,6 +108,8 @@ suite('Python envs locator - Environments Resolver', () => {
             searchLocation: Uri.file(location),
             source: [],
             type,
+            identifiedUsingNativeLocator: undefined,
+            pythonRunCommand: undefined,
         };
     }
     suite('iterEnvs()', () => {
diff --git a/src/test/pythonEnvironments/base/locators/composite/resolverUtils.unit.test.ts b/src/test/pythonEnvironments/base/locators/composite/resolverUtils.unit.test.ts
@@ -153,6 +153,8 @@ suite('Resolver Utils', () => {
                 kind: PythonEnvKind.MicrosoftStore,
                 distro: { org: 'Microsoft' },
                 source: [PythonEnvSource.PathEnvVar],
+                identifiedUsingNativeLocator: undefined,
+                pythonRunCommand: undefined,
                 ...createExpectedInterpreterInfo(python38path),
             };
             setEnvDisplayString(expected);
@@ -175,6 +177,8 @@ suite('Resolver Utils', () => {
                 kind: PythonEnvKind.MicrosoftStore,
                 distro: { org: 'Microsoft' },
                 source: [PythonEnvSource.PathEnvVar],
+                identifiedUsingNativeLocator: undefined,
+                pythonRunCommand: undefined,
                 ...createExpectedInterpreterInfo(python38path),
             };
             setEnvDisplayString(expected);
@@ -239,6 +243,8 @@ suite('Resolver Utils', () => {
                 distro: { org: '' },
                 searchLocation: undefined,
                 source: [],
+                identifiedUsingNativeLocator: undefined,
+                pythonRunCommand: undefined,
             };
             info.type = PythonEnvType.Conda;
             setEnvDisplayString(info);
@@ -351,6 +357,8 @@ suite('Resolver Utils', () => {
                 searchLocation: Uri.file(location),
                 source: [],
                 type: PythonEnvType.Virtual,
+                identifiedUsingNativeLocator: undefined,
+                pythonRunCommand: undefined,
             };
             setEnvDisplayString(info);
             return info;
@@ -406,6 +414,8 @@ suite('Resolver Utils', () => {
                 distro: { org: '' },
                 searchLocation: undefined,
                 source: [],
+                identifiedUsingNativeLocator: undefined,
+                pythonRunCommand: undefined,
             };
             setEnvDisplayString(info);
             return info;
