Allow low level locators to emit a single kind

Author: Kartik Raj

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

@@ -20,15 +20,15 @@ import {
     PythonVersion,
     virtualEnvKinds,
 } from '.';
-import { BasicEnvInfo } from '../locator';
+import { CompositeEnvInfo, convertKindIntoArray } from '../locator';
 
 /**
  * Create a new info object with all values empty.
  *
  * @param init - if provided, these values are applied to the new object
  */
 export function buildEnvInfo(init?: {
-    kind?: PythonEnvKind[];
+    kind?: PythonEnvKind[] | PythonEnvKind;
     executable?: string;
     name?: string;
     location?: string;
@@ -83,7 +83,7 @@ export function buildEnvInfo(init?: {
 export function copyEnvInfo(
     env: PythonEnvInfo,
     updates?: {
-        kind?: PythonEnvKind[];
+        kind?: PythonEnvKind[] | PythonEnvKind;
     },
 ): PythonEnvInfo {
     // We don't care whether or not extra/hidden properties
@@ -98,15 +98,15 @@ export function copyEnvInfo(
 function updateEnv(
     env: PythonEnvInfo,
     updates: {
-        kind?: PythonEnvKind[];
+        kind?: PythonEnvKind[] | PythonEnvKind;
         executable?: string;
         location?: string;
         version?: PythonVersion;
         searchLocation?: Uri;
     },
 ): void {
     if (updates.kind !== undefined) {
-        env.kind = updates.kind;
+        env.kind = convertKindIntoArray(updates.kind);
     }
     if (updates.executable !== undefined) {
         env.executable.filename = updates.executable;
@@ -173,7 +173,7 @@ function buildEnvDisplayString(env: PythonEnvInfo, getAllDetails = false): strin
  * If insufficient data is provided to generate a minimal object, such
  * that it is not identifiable, then `undefined` is returned.
  */
-function getMinimalPartialInfo(env: string | PythonEnvInfo | BasicEnvInfo): Partial<PythonEnvInfo> | undefined {
+function getMinimalPartialInfo(env: string | PythonEnvInfo | CompositeEnvInfo): Partial<PythonEnvInfo> | undefined {
     if (typeof env === 'string') {
         if (env === '') {
             return undefined;
@@ -235,8 +235,8 @@ export function getEnvID(interpreterPath: string, envFolderPath?: string): strin
  * where multiple versions of python executables are all put in the same directory.
  */
 export function areSameEnv(
-    left: string | PythonEnvInfo | BasicEnvInfo,
-    right: string | PythonEnvInfo | BasicEnvInfo,
+    left: string | PythonEnvInfo | CompositeEnvInfo,
+    right: string | PythonEnvInfo | CompositeEnvInfo,
     allowPartialMatch = true,
 ): boolean | undefined {
     const leftInfo = getMinimalPartialInfo(left);

src/client/pythonEnvironments/base/locator.ts

@@ -296,14 +296,31 @@ export type PythonLocatorQuery = BasicPythonLocatorQuery & {
 
 type QueryForEvent<E> = E extends PythonEnvsChangedEvent ? PythonLocatorQuery : BasicPythonLocatorQuery;
 
-export type BasicEnvInfo = {
-    kind: PythonEnvKind[];
+export type BasicEnvInfo<T = PythonEnvKind[] | PythonEnvKind> = {
+    kind: T;
     executablePath: string;
     source?: PythonEnvSource[];
     envPath?: string;
     extensionId?: ExtensionID;
 };
 
+/**
+ * A version of `BasicEnvInfo` used for composite locators.
+ */
+export type CompositeEnvInfo = BasicEnvInfo<PythonEnvKind[]>;
+
+export function convertBasicToComposite(env: BasicEnvInfo): CompositeEnvInfo {
+    env.kind = convertKindIntoArray(env.kind);
+    return env as CompositeEnvInfo;
+}
+
+export function convertKindIntoArray(kind: PythonEnvKind | PythonEnvKind[]): PythonEnvKind[] {
+    if (!Array.isArray(kind)) {
+        kind = [kind];
+    }
+    return kind;
+}
+
 /**
  * A single Python environment locator.
  *

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

@@ -9,6 +9,8 @@ import { areSameEnv } from '../../info/env';
 import { sortExtensionSource, sortKindFunction } from '../../info/envKind';
 import {
     BasicEnvInfo,
+    CompositeEnvInfo,
+    convertBasicToComposite,
     ILocator,
     IPythonEnvsIterator,
     isProgressEvent,
@@ -22,7 +24,7 @@ import { PythonEnvsChangedEvent } from '../../watcher';
 /**
  * Combines duplicate environments received from the incoming locator into one and passes on unique environments
  */
-export class PythonEnvsReducer implements ILocator<BasicEnvInfo> {
+export class PythonEnvsReducer implements ILocator<CompositeEnvInfo> {
     public get onChanged(): Event<PythonEnvsChangedEvent> {
         return this.parentLocator.onChanged;
     }
@@ -31,8 +33,8 @@ export class PythonEnvsReducer implements ILocator<BasicEnvInfo> {
 
     constructor(private readonly parentLocator: ILocator<BasicEnvInfo>) {}
 
-    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<BasicEnvInfo> {
-        const didUpdate = new EventEmitter<PythonEnvUpdatedEvent<BasicEnvInfo> | ProgressNotificationEvent>();
+    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<CompositeEnvInfo> {
+        const didUpdate = new EventEmitter<PythonEnvUpdatedEvent<CompositeEnvInfo> | ProgressNotificationEvent>();
         const incomingIterator = this.parentLocator.iterEnvs(query);
         const iterator = iterEnvsIterator(incomingIterator, didUpdate);
         iterator.onUpdated = didUpdate.event;
@@ -42,13 +44,13 @@ export class PythonEnvsReducer implements ILocator<BasicEnvInfo> {
 
 async function* iterEnvsIterator(
     iterator: IPythonEnvsIterator<BasicEnvInfo>,
-    didUpdate: EventEmitter<PythonEnvUpdatedEvent<BasicEnvInfo> | ProgressNotificationEvent>,
-): IPythonEnvsIterator<BasicEnvInfo> {
+    didUpdate: EventEmitter<PythonEnvUpdatedEvent<CompositeEnvInfo> | ProgressNotificationEvent>,
+): IPythonEnvsIterator<CompositeEnvInfo> {
     const state = {
         done: false,
         pending: 0,
     };
-    const seen: BasicEnvInfo[] = [];
+    const seen: CompositeEnvInfo[] = [];
 
     if (iterator.onUpdated !== undefined) {
         const listener = iterator.onUpdated((event) => {
@@ -66,8 +68,8 @@ async function* iterEnvsIterator(
                 );
             } else if (seen[event.index] !== undefined) {
                 const oldEnv = seen[event.index];
-                seen[event.index] = event.update;
-                didUpdate.fire({ index: event.index, old: oldEnv, update: event.update });
+                seen[event.index] = convertBasicToComposite(event.update);
+                didUpdate.fire({ index: event.index, old: oldEnv, update: seen[event.index] });
             } else {
                 // This implies a problem in a downstream locator
                 traceVerbose(`Expected already iterated env, got ${event.old} (#${event.index})`);
@@ -81,7 +83,7 @@ async function* iterEnvsIterator(
 
     let result = await iterator.next();
     while (!result.done) {
-        const currEnv = result.value;
+        const currEnv = convertBasicToComposite(result.value);
         const oldIndex = seen.findIndex((s) => areSameEnv(s, currEnv));
         if (oldIndex !== -1) {
             resolveDifferencesInBackground(oldIndex, currEnv, state, didUpdate, seen).ignoreErrors();
@@ -100,10 +102,10 @@ async function* iterEnvsIterator(
 
 async function resolveDifferencesInBackground(
     oldIndex: number,
-    newEnv: BasicEnvInfo,
+    newEnv: CompositeEnvInfo,
     state: { done: boolean; pending: number },
-    didUpdate: EventEmitter<PythonEnvUpdatedEvent<BasicEnvInfo> | ProgressNotificationEvent>,
-    seen: BasicEnvInfo[],
+    didUpdate: EventEmitter<PythonEnvUpdatedEvent<CompositeEnvInfo> | ProgressNotificationEvent>,
+    seen: CompositeEnvInfo[],
 ) {
     state.pending += 1;
     // It's essential we increment the pending call count before any asynchronus calls in this method.
@@ -125,15 +127,15 @@ async function resolveDifferencesInBackground(
  */
 function checkIfFinishedAndNotify(
     state: { done: boolean; pending: number },
-    didUpdate: EventEmitter<PythonEnvUpdatedEvent<BasicEnvInfo> | ProgressNotificationEvent>,
+    didUpdate: EventEmitter<PythonEnvUpdatedEvent<CompositeEnvInfo> | ProgressNotificationEvent>,
 ) {
     if (state.done && state.pending === 0) {
         didUpdate.fire({ stage: ProgressReportStage.discoveryFinished });
         didUpdate.dispose();
     }
 }
 
-function resolveEnvCollision(oldEnv: BasicEnvInfo, newEnv: BasicEnvInfo): BasicEnvInfo {
+function resolveEnvCollision(oldEnv: CompositeEnvInfo, newEnv: CompositeEnvInfo): CompositeEnvInfo {
     const [env] = sortEnvInfoByPriority(oldEnv, newEnv);
     const merged = cloneDeep(env);
     merged.source = union(oldEnv.source ?? [], newEnv.source ?? []);
@@ -145,8 +147,8 @@ function resolveEnvCollision(oldEnv: BasicEnvInfo, newEnv: BasicEnvInfo): BasicE
  * Selects an environment based on the environment selection priority. This should
  * match the priority in the environment identifier.
  */
-function sortEnvInfoByPriority(...envs: BasicEnvInfo[]): BasicEnvInfo[] {
-    return envs.sort((a: BasicEnvInfo, b: BasicEnvInfo) => {
+function sortEnvInfoByPriority(...envs: CompositeEnvInfo[]): CompositeEnvInfo[] {
+    return envs.sort((a: CompositeEnvInfo, b: CompositeEnvInfo) => {
         const kindDiff = sortKindFunction(getTopKind(a.kind), getTopKind(b.kind));
         if (kindDiff !== 0) {
             return kindDiff;

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

@@ -9,7 +9,7 @@ import { PythonEnvInfo, PythonEnvKind } from '../../info';
 import { getEnvPath, setEnvDisplayString } from '../../info/env';
 import { InterpreterInformation } from '../../info/interpreter';
 import {
-    BasicEnvInfo,
+    CompositeEnvInfo,
     IInternalEnvironmentProvider,
     ILocator,
     InternalEnvironmentProviderMetadata,
@@ -22,7 +22,7 @@ import {
     PythonLocatorQuery,
 } from '../../locator';
 import { PythonEnvsChangedEvent } from '../../watcher';
-import { registerResolver, resolveBasicEnv } from './resolverUtils';
+import { registerResolver, resolveCompositeEnv } from './resolverUtils';
 import { traceVerbose, traceWarn } from '../../../../logging';
 import { getEnvironmentDirFromPath, getInterpreterPathFromDir, isPythonExecutable } from '../../../common/commonUtils';
 import { getEmptyVersion } from '../../info/pythonVersion';
@@ -58,15 +58,15 @@ export class PythonEnvsResolver implements IResolvingLocator {
     }
 
     constructor(
-        private readonly parentLocator: ILocator<BasicEnvInfo>,
+        private readonly parentLocator: ILocator<CompositeEnvInfo>,
         private readonly environmentInfoService: IEnvironmentInfoService,
     ) {}
 
     public async resolveEnv(path: string): Promise<PythonEnvInfo | undefined> {
         const [executablePath, envPath] = await getExecutablePathAndEnvPath(path);
         path = executablePath.length ? executablePath : envPath;
         const kind = await identifyEnvironment(path);
-        const environment = await resolveBasicEnv({ kind: [kind], executablePath, envPath });
+        const environment = await resolveCompositeEnv({ kind: [kind], executablePath, envPath });
         if (!environment) {
             return undefined;
         }
@@ -86,7 +86,7 @@ export class PythonEnvsResolver implements IResolvingLocator {
     }
 
     private async *iterEnvsIterator(
-        iterator: IPythonEnvsIterator<BasicEnvInfo>,
+        iterator: IPythonEnvsIterator<CompositeEnvInfo>,
         didUpdate: EventEmitter<PythonEnvUpdatedEvent | ProgressNotificationEvent>,
     ): IPythonEnvsIterator {
         const state = {
@@ -112,7 +112,7 @@ export class PythonEnvsResolver implements IResolvingLocator {
                     );
                 } else if (seen[event.index] !== undefined) {
                     const old = seen[event.index];
-                    const env = await resolveBasicEnv(event.update, true);
+                    const env = await resolveCompositeEnv(event.update, true);
                     didUpdate.fire({ old, index: event.index, update: env });
                     if (env) {
                         seen[event.index] = env;
@@ -132,7 +132,7 @@ export class PythonEnvsResolver implements IResolvingLocator {
         let result = await iterator.next();
         while (!result.done) {
             // Use cache from the current refresh where possible.
-            const currEnv = await resolveBasicEnv(result.value, true);
+            const currEnv = await resolveCompositeEnv(result.value, true);
             if (currEnv) {
                 seen.push(currEnv);
                 yield currEnv;

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

@@ -23,12 +23,12 @@ import { getPyenvVersionsDir, parsePyenvVersion } from '../../../common/environm
 import { Architecture, getOSType, OSType } from '../../../../common/utils/platform';
 import { getPythonVersionFromPath as parsePythonVersionFromPath, parseVersion } from '../../info/pythonVersion';
 import { getRegistryInterpreters, getRegistryInterpretersSync } from '../../../common/windowsUtils';
-import { BasicEnvInfo } from '../../locator';
+import { CompositeEnvInfo } from '../../locator';
 import { parseVersionFromExecutable } from '../../info/executable';
 import { traceError, traceWarn } from '../../../../logging';
 import { sortExtensionSource } from '../../info/envKind';
 
-type ResolverType = (_: BasicEnvInfo, useCache?: boolean) => Promise<PythonEnvInfo | undefined>;
+type ResolverType = (_: CompositeEnvInfo, useCache?: boolean) => Promise<PythonEnvInfo | undefined>;
 type ResolversType = { resolver: ResolverType; extensionId?: string }[];
 const resolvers = new Map<PythonEnvKind, ResolverType | ResolversType>();
 Object.values(PythonEnvKind).forEach((k) => {
@@ -43,7 +43,7 @@ resolvers.set(PythonEnvKind.Pyenv, resolvePyenvEnv);
 
 export function registerResolver(
     kind: PythonEnvKind,
-    resolver: (_: BasicEnvInfo, useCache?: boolean) => Promise<PythonEnvInfo | undefined>,
+    resolver: (_: CompositeEnvInfo, useCache?: boolean) => Promise<PythonEnvInfo | undefined>,
     extensionId: string,
 ): void {
     const resolversForKind = resolvers.get(kind);
@@ -59,16 +59,19 @@ export function registerResolver(
 }
 
 /**
- * Find as much info about the given Basic Python env as possible without running the
+ * Find as much info about the given Python env as possible without running the
  * executable and returns it. Notice `undefined` is never returned, so environment
  * returned could still be invalid.
  */
-export async function resolveBasicEnv(env: BasicEnvInfo, useCache = false): Promise<PythonEnvInfo | undefined> {
+export async function resolveCompositeEnv(env: CompositeEnvInfo, useCache = false): Promise<PythonEnvInfo | undefined> {
     const { kind, source } = env;
-    const value = resolvers.get(kind[0]);
+    let value = resolvers.get(kind[0]);
     if (!value) {
-        traceError('No resolver found for kind:', kind[0]);
-        return undefined;
+        value = env.extensionId ? resolvers.get(env.extensionId as PythonEnvKind) : undefined;
+        if (!value) {
+            traceError('No resolver found for env:', JSON.stringify(env));
+            return undefined;
+        }
     }
     let resolverForKind: ResolverType;
     if (Array.isArray(value)) {
@@ -146,7 +149,7 @@ async function updateEnvUsingRegistry(env: PythonEnvInfo): Promise<void> {
     }
 }
 
-async function resolveGloballyInstalledEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
+async function resolveGloballyInstalledEnv(env: CompositeEnvInfo): Promise<PythonEnvInfo> {
     const { executablePath } = env;
     let version;
     try {
@@ -162,7 +165,7 @@ async function resolveGloballyInstalledEnv(env: BasicEnvInfo): Promise<PythonEnv
     return envInfo;
 }
 
-async function resolveSimpleEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
+async function resolveSimpleEnv(env: CompositeEnvInfo): Promise<PythonEnvInfo> {
     const { executablePath, kind } = env;
     const envInfo = buildEnvInfo({
         kind,
@@ -175,7 +178,7 @@ async function resolveSimpleEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
     return envInfo;
 }
 
-async function resolveCondaEnv(env: BasicEnvInfo, useCache?: boolean): Promise<PythonEnvInfo> {
+async function resolveCondaEnv(env: CompositeEnvInfo, useCache?: boolean): Promise<PythonEnvInfo> {
     const { executablePath } = env;
     const conda = await Conda.getConda();
     if (conda === undefined) {
@@ -184,7 +187,7 @@ async function resolveCondaEnv(env: BasicEnvInfo, useCache?: boolean): Promise<P
     const envs = (await conda?.getEnvList(useCache)) ?? [];
     for (const { name, prefix } of envs) {
         let executable = await getInterpreterPathFromDir(prefix);
-        const currEnv: BasicEnvInfo = { executablePath: executable ?? '', kind: env.kind, envPath: prefix };
+        const currEnv: CompositeEnvInfo = { executablePath: executable ?? '', kind: env.kind, envPath: prefix };
         if (areSameEnv(env, currEnv)) {
             if (env.executablePath.length > 0) {
                 executable = env.executablePath;
@@ -215,7 +218,7 @@ async function resolveCondaEnv(env: BasicEnvInfo, useCache?: boolean): Promise<P
     return resolveSimpleEnv(env);
 }
 
-async function resolvePyenvEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
+async function resolvePyenvEnv(env: CompositeEnvInfo): Promise<PythonEnvInfo> {
     const { executablePath } = env;
     const location = getEnvironmentDirFromPath(executablePath);
     const name = path.basename(location);
@@ -268,7 +271,7 @@ async function isBaseCondaPyenvEnvironment(executablePath: string) {
     return arePathsSame(path.dirname(location), pyenvVersionDir);
 }
 
-async function resolveWindowsStoreEnv(env: BasicEnvInfo): Promise<PythonEnvInfo> {
+async function resolveWindowsStoreEnv(env: CompositeEnvInfo): Promise<PythonEnvInfo> {
     const { executablePath } = env;
     return buildEnvInfo({
         kind: env.kind,

src/client/pythonEnvironments/index.ts

@@ -106,7 +106,7 @@ async function createLocator(
     // This is shared.
 ): Promise<IDiscoveryAPI> {
     // Create the low-level locators.
-    let locators: ILocator<BasicEnvInfo> = new ExtensionLocators(
+    const locators: ILocator<BasicEnvInfo> = new ExtensionLocators(
         // Here we pull the locators together.
         createNonWorkspaceLocators(ext),
         createWorkspaceLocator(ext),
@@ -116,9 +116,9 @@ async function createLocator(
     const envInfoService = getEnvironmentInfoService(ext.disposables);
 
     // Build the stack of composite locators.
-    locators = new PythonEnvsReducer(locators);
+    const reducer = new PythonEnvsReducer(locators);
     const resolvingLocator = new PythonEnvsResolver(
-        locators,
+        reducer,
         // These are shared.
         envInfoService,
     );