Add environments reducer

Author: Kartik Raj

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,22 @@ 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,
+    environment2: PythonEnvInfo,
+): boolean {
+    if (!environment1 || !environment2) {
+        return false;
+    }
+    if (arePathsSame(environment1.executable.filename, environment2.executable.filename)) {
+        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,127 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { Event, EventEmitter } from 'vscode';
+import { areSameEnvironment, PythonEnvInfo, PythonEnvKind } from '../base/info';
+import {
+    ILocator, IPythonEnvsIterator, PythonEnvUpdatedEvent, QueryForEvent,
+} from '../base/locator';
+import { PythonEnvsChangedEvent } from '../base/watcher';
+
+export class PythonEnvsReducer implements ILocator {
+    public get onChanged(): Event<PythonEnvsChangedEvent> {
+        return this.pythonEnvsManager.onChanged;
+    }
+
+    constructor(private readonly pythonEnvsManager: ILocator) {}
+
+    public resolveEnv(env: string | PythonEnvInfo): Promise<PythonEnvInfo | undefined> {
+        return this.pythonEnvsManager.resolveEnv(env);
+    }
+
+    public iterEnvs(query?: QueryForEvent<PythonEnvsChangedEvent>): IPythonEnvsIterator {
+        const didUpdate = new EventEmitter<PythonEnvUpdatedEvent | null>();
+        const iterator: IPythonEnvsIterator = this.iterEnvsIterator(didUpdate, query);
+        iterator.onUpdated = didUpdate.event;
+        return iterator;
+    }
+
+    private async* iterEnvsIterator(
+        didUpdate: EventEmitter<PythonEnvUpdatedEvent | null>,
+        query?: QueryForEvent<PythonEnvsChangedEvent>,
+    ): AsyncIterator<PythonEnvInfo, void> {
+        const state = {
+            done: false,
+            pending: 0,
+        };
+        const seen: PythonEnvInfo[] = [];
+        const iterator = this.pythonEnvsManager.iterEnvs(query);
+
+        if (iterator.onUpdated !== undefined) {
+            iterator.onUpdated((event) => {
+                if (event === null) {
+                    state.done = true;
+                    checkIfFinishedAndNotify(state, didUpdate);
+                } else {
+                    const old = seen.find((s) => areSameEnvironment(s, event.old));
+                    if (old !== undefined) {
+                        state.pending += 1;
+                        resolveDifferencesInBackground(old, event.new, { seen, ...state }, didUpdate).ignoreErrors();
+                    }
+                }
+            });
+        }
+
+        let result = await iterator.next();
+        while (!result.done) {
+            const currEnv = result.value;
+            const old = seen.find((s) => areSameEnvironment(s, currEnv));
+            if (old !== undefined) {
+                state.pending += 1;
+                resolveDifferencesInBackground(old, currEnv, { seen, ...state }, didUpdate).ignoreErrors();
+            } else {
+                yield currEnv;
+                seen.push(currEnv);
+            }
+            // eslint-disable-next-line no-await-in-loop
+            result = await iterator.next();
+        }
+        if (iterator.onUpdated === undefined) {
+            state.done = true;
+        }
+    }
+}
+
+async function resolveDifferencesInBackground(
+    oldEnv: PythonEnvInfo,
+    newEnv: PythonEnvInfo,
+    state: { seen: PythonEnvInfo[]; done: boolean; pending: number },
+    didUpdate: EventEmitter<PythonEnvUpdatedEvent | null>,
+) {
+    const merged = mergeEnvironments(oldEnv, newEnv);
+    didUpdate.fire({ old: oldEnv, new: merged });
+    state.pending -= 1;
+    state.seen[state.seen.indexOf(oldEnv)] = merged;
+    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 {
+    // 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) {
+        environment.kind = other.kind;
+    }
+    const props: (keyof PythonEnvInfo)[] = [
+        'version',
+        'kind',
+        'executable',
+        'name',
+        'arch',
+        'distro',
+        'defaultDisplayName',
+        'searchLocation',
+    ];
+    props.forEach((prop) => {
+        if (!environment[prop] && other[prop]) {
+            // tslint:disable: no-any
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            (environment as any)[prop] = other[prop];
+        }
+    });
+    return environment;
+}

src/test/pythonEnvironments/collection/environmentsReducer.unit.test.ts

@@ -0,0 +1,29 @@
+import { assert } from 'chai';
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { PythonEnvKind } from '../../../client/pythonEnvironments/base/info';
+import { PythonEnvsReducer } from '../../../client/pythonEnvironments/collection/environmentsReducer';
+import {
+    createLocatedEnv, getEnvs, SimpleLocator,
+} from '../base/common';
+
+suite('Environments Reducer', () => {
+    test('Duplicated incoming environments from locator manager are removed', async () => {
+        const env1 = createLocatedEnv('path/to/env1', '3.5.12b1', PythonEnvKind.Venv);
+        const env2 = createLocatedEnv('path/to/env2', '3.8.1', PythonEnvKind.Conda);
+        const env3 = createLocatedEnv('path/to/env3', '2.7', PythonEnvKind.System);
+        const env4 = createLocatedEnv('path/to/env2', '3.9.0rc2', PythonEnvKind.Pyenv);
+        const env5 = createLocatedEnv('path/to/env1', '3.8', PythonEnvKind.System);
+        const environments = [env1, env2, env3, env4, env5];
+        const pythonEnvManager = new SimpleLocator(environments);
+
+        const reducer = new PythonEnvsReducer(pythonEnvManager);
+
+        const iterator = reducer.iterEnvs();
+        const envs = await getEnvs(iterator);
+
+        const expected = [env1, env2, env3];
+        assert.deepEqual(envs.sort(), expected.sort());
+    });
+});