Added environments reducer

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

@@ -4,6 +4,7 @@
 import { Uri } from 'vscode';
 import { Architecture } from '../../../common/utils/platform';
 import { BasicVersionInfo, VersionInfo } from '../../../common/utils/version';
+import { arePathsSame } from '../../common/externalDependencies';
 
 /**
  * IDs for the various supported Python environments.
@@ -143,3 +144,31 @@ export type PythonEnvInfo = _PythonEnvInfo & {
     defaultDisplayName?: string;
     searchLocation?: Uri;
 };
+
+/**
+ * Determine if the given infos correspond to the same env.
+ *
+ * @param environment1 - one of the two envs to compare
+ * @param environment2 - one of the two envs to compare
+ */
+export function areSameEnvironment(
+    environment1: PythonEnvInfo | string,
+    environment2: PythonEnvInfo | string,
+): boolean {
+    let path1: string;
+    let path2: string;
+    if (typeof environment1 === 'string') {
+        path1 = environment1;
+    } else {
+        path1 = environment1.executable.filename;
+    }
+    if (typeof environment2 === 'string') {
+        path2 = environment2;
+    } else {
+        path2 = environment2.executable.filename;
+    }
+    if (arePathsSame(path1, path2)) {
+        return true;
+    }
+    return false;
+}

src/client/pythonEnvironments/base/locator.ts

@@ -92,7 +92,7 @@ export type PythonLocatorQuery = BasicPythonLocatorQuery & {
     searchLocations?: Uri[];
 };
 
-type QueryForEvent<E> = E extends PythonEnvsChangedEvent ? PythonLocatorQuery : BasicPythonLocatorQuery;
+export type QueryForEvent<E> = E extends PythonEnvsChangedEvent ? PythonLocatorQuery : BasicPythonLocatorQuery;
 
 /**
  * A single Python environment locator.

src/client/pythonEnvironments/collection/environmentsReducer.ts

@@ -0,0 +1,165 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { cloneDeep, isEqual } from 'lodash';
+import { Event, EventEmitter } from 'vscode';
+import { traceVerbose } from '../../common/logger';
+import { createDeferred } from '../../common/utils/async';
+import { areSameEnvironment, PythonEnvInfo, PythonEnvKind } from '../base/info';
+import {
+    ILocator, IPythonEnvsIterator, PythonEnvUpdatedEvent, QueryForEvent,
+} from '../base/locator';
+import { PythonEnvsChangedEvent } from '../base/watcher';
+
+/**
+ * Combines duplicate environments received from the incoming locator into one and passes on unique environments
+ */
+export class PythonEnvsReducer implements ILocator {
+    public get onChanged(): Event<PythonEnvsChangedEvent> {
+        return this.pythonEnvsManager.onChanged;
+    }
+
+    constructor(private readonly pythonEnvsManager: ILocator) {}
+
+    public async resolveEnv(env: string | PythonEnvInfo): Promise<PythonEnvInfo | undefined> {
+        let environment: PythonEnvInfo | undefined;
+        const waitForUpdatesDeferred = createDeferred<void>();
+        const iterator = this.iterEnvs();
+        iterator.onUpdated!((event) => {
+            if (event === null) {
+                waitForUpdatesDeferred.resolve();
+            } else if (environment && areSameEnvironment(environment, event.new)) {
+                environment = event.new;
+            }
+        });
+        let result = await iterator.next();
+        while (!result.done) {
+            if (areSameEnvironment(result.value, env)) {
+                environment = result.value;
+            }
+            // eslint-disable-next-line no-await-in-loop
+            result = await iterator.next();
+        }
+        if (!environment) {
+            return undefined;
+        }
+        await waitForUpdatesDeferred.promise;
+        return this.pythonEnvsManager.resolveEnv(environment);
+    }
+
+    public iterEnvs(query?: QueryForEvent<PythonEnvsChangedEvent>): IPythonEnvsIterator {
+        const didUpdate = new EventEmitter<PythonEnvUpdatedEvent | null>();
+        const incomingIterator = this.pythonEnvsManager.iterEnvs(query);
+        const iterator: IPythonEnvsIterator = iterEnvsIterator(incomingIterator, didUpdate);
+        iterator.onUpdated = didUpdate.event;
+        return iterator;
+    }
+}
+
+async function* iterEnvsIterator(
+    iterator: IPythonEnvsIterator,
+    didUpdate: EventEmitter<PythonEnvUpdatedEvent | null>,
+): AsyncIterator<PythonEnvInfo, void> {
+    const state = {
+        done: false,
+        pending: 0,
+    };
+    const seen: PythonEnvInfo[] = [];
+
+    if (iterator.onUpdated !== undefined) {
+        iterator.onUpdated((event) => {
+            if (event === null) {
+                state.done = true;
+                checkIfFinishedAndNotify(state, didUpdate);
+            } else {
+                const oldIndex = seen.findIndex((s) => areSameEnvironment(s, event.old));
+                if (oldIndex !== -1) {
+                    state.pending += 1;
+                    resolveDifferencesInBackground(oldIndex, event.new, state, didUpdate, seen).ignoreErrors();
+                } else {
+                    // This implies a problem in a downstream locator
+                    traceVerbose(`Expected already iterated env, got ${event.old}`);
+                }
+            }
+        });
+    }
+
+    let result = await iterator.next();
+    while (!result.done) {
+        const currEnv = result.value;
+        const oldIndex = seen.findIndex((s) => areSameEnvironment(s, currEnv));
+        if (oldIndex !== -1) {
+            state.pending += 1;
+            resolveDifferencesInBackground(oldIndex, currEnv, state, didUpdate, seen).ignoreErrors();
+        } else {
+            // We haven't yielded a matching env so yield this one as-is.
+            yield currEnv;
+            seen.push(currEnv);
+        }
+        // eslint-disable-next-line no-await-in-loop
+        result = await iterator.next();
+    }
+    if (iterator.onUpdated === undefined) {
+        state.done = true;
+        checkIfFinishedAndNotify(state, didUpdate);
+    }
+}
+
+async function resolveDifferencesInBackground(
+    oldIndex: number,
+    newEnv: PythonEnvInfo,
+    state: { done: boolean; pending: number },
+    didUpdate: EventEmitter<PythonEnvUpdatedEvent | null>,
+    seen: PythonEnvInfo[],
+) {
+    const oldEnv = seen[oldIndex];
+    const merged = mergeEnvironments(oldEnv, newEnv);
+    if (!isEqual(oldEnv, merged)) {
+        didUpdate.fire({ old: oldEnv, new: merged });
+        seen[oldIndex] = merged;
+    }
+    state.pending -= 1;
+    checkIfFinishedAndNotify(state, didUpdate);
+}
+
+/**
+ * When all info from incoming iterator has been received and all background calls finishes, notify that we're done
+ * @param state Carries the current state of progress
+ * @param didUpdate Used to notify when finished
+ */
+function checkIfFinishedAndNotify(
+    state: { done: boolean; pending: number },
+    didUpdate: EventEmitter<PythonEnvUpdatedEvent | null>,
+) {
+    if (state.done && state.pending === 0) {
+        didUpdate.fire(null);
+        didUpdate.dispose();
+    }
+}
+
+export function mergeEnvironments(environment: PythonEnvInfo, other: PythonEnvInfo): PythonEnvInfo {
+    const result = cloneDeep(environment);
+    // Preserve type information.
+    // Possible we identified environment as unknown, but a later provider has identified env type.
+    if (environment.kind === PythonEnvKind.Unknown && other.kind && other.kind !== PythonEnvKind.Unknown) {
+        result.kind = other.kind;
+    }
+    const props: (keyof PythonEnvInfo)[] = [
+        'version',
+        'kind',
+        'executable',
+        'name',
+        'arch',
+        'distro',
+        'defaultDisplayName',
+        'searchLocation',
+    ];
+    props.forEach((prop) => {
+        if (!result[prop] && other[prop]) {
+            // tslint:disable: no-any
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            (result as any)[prop] = other[prop];
+        }
+    });
+    return result;
+}

src/test/pythonEnvironments/base/common.ts

@@ -1,14 +1,19 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { createDeferred, flattenIterator, iterable, mapToIterator } from '../../../client/common/utils/async';
+import { Event } from 'vscode';
+import {
+    createDeferred, flattenIterator, iterable, mapToIterator,
+} from '../../../client/common/utils/async';
 import { Architecture } from '../../../client/common/utils/platform';
 import {
     PythonEnvInfo,
     PythonEnvKind,
 } from '../../../client/pythonEnvironments/base/info';
 import { parseVersion } from '../../../client/pythonEnvironments/base/info/pythonVersion';
-import { IPythonEnvsIterator, Locator, PythonLocatorQuery } from '../../../client/pythonEnvironments/base/locator';
+import {
+    IPythonEnvsIterator, Locator, PythonEnvUpdatedEvent, PythonLocatorQuery,
+} from '../../../client/pythonEnvironments/base/locator';
 import { PythonEnvsChangedEvent } from '../../../client/pythonEnvironments/base/watcher';
 
 export function createEnv(
@@ -66,6 +71,7 @@ export class SimpleLocator extends Locator {
             resolve?: null | ((env: PythonEnvInfo) => Promise<PythonEnvInfo | undefined>);
             before?: Promise<void>;
             after?: Promise<void>;
+            onUpdated?: Event<PythonEnvUpdatedEvent | null>;
             beforeEach?(e: PythonEnvInfo): Promise<void>;
             afterEach?(e: PythonEnvInfo): Promise<void>;
             onQuery?(query: PythonLocatorQuery | undefined, envs: PythonEnvInfo[]): Promise<PythonEnvInfo[]>;
@@ -83,7 +89,7 @@ export class SimpleLocator extends Locator {
         const deferred = this.deferred;
         const callbacks = this.callbacks;
         let envs = this.envs;
-        async function* iterator() {
+        const iterator: IPythonEnvsIterator = async function*() {
             if (callbacks?.onQuery !== undefined) {
                 envs = await callbacks.onQuery(query, envs);
             }
@@ -114,8 +120,9 @@ export class SimpleLocator extends Locator {
                 await callbacks.after;
             }
             deferred.resolve();
-        }
-        return iterator();
+        }();
+        iterator.onUpdated = this.callbacks?.onUpdated;
+        return iterator;
     }
     public async resolveEnv(env: string | PythonEnvInfo): Promise<PythonEnvInfo | undefined> {
         const envInfo: PythonEnvInfo = typeof env === 'string' ? createEnv('', '', undefined, env) : env;