Added environments resolver

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/environmentsResolver.ts

@@ -0,0 +1,126 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { cloneDeep } from 'lodash';
+import { Event, EventEmitter } from 'vscode';
+import { traceVerbose } from '../../common/logger';
+import { areSameEnvironment, PythonEnvInfo } from '../base/info';
+import { InterpreterInformation } from '../base/info/interpreter';
+import {
+    ILocator, IPythonEnvsIterator, PythonEnvUpdatedEvent, QueryForEvent,
+} from '../base/locator';
+import { PythonEnvsChangedEvent } from '../base/watcher';
+import { IEnvironmentInfoService } from '../info/environmentInfoService';
+
+export class PythonEnvsReducer implements ILocator {
+    public get onChanged(): Event<PythonEnvsChangedEvent> {
+        return this.pythonEnvsReducer.onChanged;
+    }
+
+    constructor(
+        private readonly pythonEnvsReducer: ILocator,
+        private readonly environmentInfoService: IEnvironmentInfoService,
+    ) {}
+
+    public resolveEnv(env: string | PythonEnvInfo): Promise<PythonEnvInfo | undefined> {
+        return this.pythonEnvsReducer.resolveEnv(env);
+    }
+
+    public iterEnvs(query?: QueryForEvent<PythonEnvsChangedEvent>): IPythonEnvsIterator {
+        const didUpdate = new EventEmitter<PythonEnvUpdatedEvent | null>();
+        const incomingIterator = this.pythonEnvsReducer.iterEnvs(query);
+        const iterator: IPythonEnvsIterator = this.iterEnvsIterator(incomingIterator, didUpdate);
+        iterator.onUpdated = didUpdate.event;
+        return iterator;
+    }
+
+    private async* 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 old = seen.find((s) => areSameEnvironment(s, event.old));
+                    if (old !== undefined) {
+                        didUpdate.fire({ old, new: event.new });
+                        seen[seen.indexOf(old)] = event.new;
+                        state.pending += 1;
+                        this.resolveInBackground(event.new, state, didUpdate, seen).ignoreErrors();
+                    } else {
+                        // This implies a problem in a downstream locator
+                        traceVerbose(`Expected already iterated env in resolver, got ${event.old}`);
+                    }
+                }
+            });
+        }
+
+        let result = await iterator.next();
+        while (!result.done) {
+            const currEnv = result.value;
+            // Resolver only expects unique environments, so store & yield as-is.
+            seen.push(currEnv);
+            yield currEnv;
+            state.pending += 1;
+            this.resolveInBackground(currEnv, state, didUpdate, seen).ignoreErrors();
+            // eslint-disable-next-line no-await-in-loop
+            result = await iterator.next();
+        }
+        if (iterator.onUpdated === undefined) {
+            state.done = true;
+            checkIfFinishedAndNotify(state, didUpdate);
+        }
+    }
+
+    private async resolveInBackground(
+        env: PythonEnvInfo,
+        state: { done: boolean; pending: number },
+        didUpdate: EventEmitter<PythonEnvUpdatedEvent | null>,
+        seen: PythonEnvInfo[],
+    ) {
+        const interpreterInfo = await this.environmentInfoService.getEnvironmentInfo(env.executable.filename);
+        // Old environment to be fired might have changed until now, search for it again
+        const seenEnv = seen.find((s) => areSameEnvironment(s, env));
+        if (interpreterInfo && seenEnv) {
+            const resolvedEnv = getResolvedEnv(interpreterInfo, seenEnv);
+            didUpdate.fire({ old: seenEnv, new: resolvedEnv });
+            seen[seen.indexOf(seenEnv)] = resolvedEnv;
+        }
+        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();
+    }
+}
+
+function getResolvedEnv(interpreterInfo: InterpreterInformation, environment: PythonEnvInfo) {
+    // Deep copy into a new object
+    const resolvedEnv = cloneDeep(environment);
+    resolvedEnv.version = interpreterInfo.version;
+    resolvedEnv.executable.filename = interpreterInfo.executable.filename;
+    resolvedEnv.executable.sysPrefix = interpreterInfo.executable.sysPrefix;
+    resolvedEnv.arch = interpreterInfo.arch;
+    return resolvedEnv;
+}