microsoft · kimadeline · Jul 14, 2021 · Jun 23, 2021 · Jun 23, 2021 · Jun 23, 2021
diff --git a/src/client/interpreter/autoSelection/index.ts b/src/client/interpreter/autoSelection/index.ts
@@ -6,15 +6,18 @@
 import { inject, injectable, named } from 'inversify';
 import { Event, EventEmitter, Uri } from 'vscode';
 import { IWorkspaceService } from '../../common/application/types';
+import { EnvironmentSorting } from '../../common/experiments/groups';
 import '../../common/extensions';
 import { IFileSystem } from '../../common/platform/types';
-import { IPersistentState, IPersistentStateFactory, Resource } from '../../common/types';
+import { IExperimentService, IPersistentState, IPersistentStateFactory, Resource } from '../../common/types';
 import { createDeferred, Deferred } from '../../common/utils/async';
 import { compareSemVerLikeVersions } from '../../pythonEnvironments/base/info/pythonVersion';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { captureTelemetry, sendTelemetryEvent } from '../../telemetry';
 import { EventName } from '../../telemetry/constants';
-import { IInterpreterHelper } from '../contracts';
+import { EnvTypeHeuristic, getEnvTypeHeuristic } from '../configuration/environmentTypeComparer';
+import { InterpreterComparisonType, IInterpreterComparer } from '../configuration/types';
+import { IInterpreterHelper, IInterpreterService } from '../contracts';
 import {
     AutoSelectionRule,
     IInterpreterAutoSelectionRule,
@@ -46,6 +49,11 @@ export class InterpreterAutoSelectionService implements IInterpreterAutoSelectio
         @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService,
         @inject(IPersistentStateFactory) private readonly stateFactory: IPersistentStateFactory,
         @inject(IFileSystem) private readonly fs: IFileSystem,
+        @inject(IExperimentService) private readonly experimentService: IExperimentService,
+        @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
+        @inject(IInterpreterComparer)
+        @named(InterpreterComparisonType.EnvType)
+        private readonly envTypeComparer: IInterpreterComparer,
         @inject(IInterpreterAutoSelectionRule)
         @named(AutoSelectionRule.systemWide)
         systemInterpreter: IInterpreterAutoSelectionRule,
@@ -104,19 +112,30 @@ export class InterpreterAutoSelectionService implements IInterpreterAutoSelectio
         winRegInterpreter.setNextRule(systemInterpreter);
     }
 
+    /**
+     * If there's a cached auto-selected interpreter -> return it.
+     * If not, check if we are in the env sorting experiment, and use the appropriate auto-selection logic.
+     */
     @captureTelemetry(EventName.PYTHON_INTERPRETER_AUTO_SELECTION, { rule: AutoSelectionRule.all }, true)
     public async autoSelectInterpreter(resource: Resource): Promise<void> {
         const key = this.getWorkspacePathKey(resource);
+
         if (!this.autoSelectedWorkspacePromises.has(key)) {
             const deferred = createDeferred<void>();
             this.autoSelectedWorkspacePromises.set(key, deferred);
+
             await this.initializeStore(resource);
             await this.clearWorkspaceStoreIfInvalid(resource);
-            await this.userDefinedInterpreter.autoSelectInterpreter(resource, this);
-            this.didAutoSelectedInterpreterEmitter.fire();
-            Promise.all(this.rules.map((item) => item.autoSelectInterpreter(resource))).ignoreErrors();
+
+            if (await this.experimentService.inExperiment(EnvironmentSorting.experiment)) {
+                await this.autoselectInterpreterWithLocators(resource);
+            } else {
+                await this.autoselectInterpreterWithRules(resource);
+            }
+
             deferred.resolve();
         }
+
         return this.autoSelectedWorkspacePromises.get(key)!.promise;
     }
 
@@ -221,4 +240,44 @@ export class InterpreterAutoSelectionService implements IInterpreterAutoSelectio
         }
         return undefined;
     }
+
+    private async autoselectInterpreterWithRules(resource: Resource): Promise<void> {
+        await this.userDefinedInterpreter.autoSelectInterpreter(resource, this);
+
+        this.didAutoSelectedInterpreterEmitter.fire();
+
+        Promise.all(this.rules.map((item) => item.autoSelectInterpreter(resource))).ignoreErrors();
+    }
+
+    /**
+     * Auto-selection logic:
+     * 1. If there are cached interpreters (not the first session in this workspace)
+     *      -> sort using the same logic as in the interpreter quickpick and return the first one;
+     * 2. If not, we already fire all the locators, so wait for their response, sort the interpreters and return the first one.
+     *
+     * `getInterpreters` will check the cache first and return early if there are any cached interpreters,
+     * and if not it will wait for locators to return.
+     * As such, we can sort interpreters based on what it returns.
+     */
+    private async autoselectInterpreterWithLocators(resource: Resource): Promise<void> {
+        const interpreters = await this.interpreterService.getInterpreters(resource);
+        const workspaceUri = this.interpreterHelper.getActiveWorkspaceUri(resource);
+
+        // When auto-selecting an intepreter for a workspace, we either want to return a local one
+        // or fallback on a globally-installed interpreter, and we don't want want to suggest a global environment
+        // because we would have to add a way to match environments to a workspace.
+        const filteredInterpreters = interpreters.filter(
+            (i) => getEnvTypeHeuristic(i, workspaceUri?.folderUri.fsPath || '') !== EnvTypeHeuristic.Global,
+        );
+
+        filteredInterpreters.sort(this.envTypeComparer.compare.bind(this.envTypeComparer));
+
+        if (workspaceUri) {
+            this.setWorkspaceInterpreter(workspaceUri.folderUri, filteredInterpreters[0]);
+        } else {
+            this.setGlobalInterpreter(filteredInterpreters[0]);
+        }
+
+        this.didAutoSelectedInterpreterEmitter.fire();
+    }
 }
diff --git a/src/client/interpreter/configuration/environmentTypeComparer.ts b/src/client/interpreter/configuration/environmentTypeComparer.ts
@@ -15,7 +15,7 @@ import { IInterpreterComparer } from './types';
  * - Global environments (pipenv, conda);
  * - Globally-installed interpreters (/usr/bin/python3, Windows Store).
  */
-enum EnvTypeHeuristic {
+export enum EnvTypeHeuristic {
     Local = 1,
     Global = 2,
     GlobalInterpreters = 3,
@@ -165,7 +165,7 @@ function compareEnvironmentType(a: PythonEnvironment, b: PythonEnvironment, work
 /**
  * Return a heuristic value depending on the environment type.
  */
-function getEnvTypeHeuristic(environment: PythonEnvironment, workspacePath: string): EnvTypeHeuristic {
+export function getEnvTypeHeuristic(environment: PythonEnvironment, workspacePath: string): EnvTypeHeuristic {
     const { envType } = environment;
 
     if (workspacePath.length > 0 && environment.envPath && isParentPath(environment.envPath, workspacePath)) {

diff --git a/src/test/configuration/environmentTypeComparer.unit.test.ts b/src/test/configuration/environmentTypeComparer.unit.test.ts
@@ -4,7 +4,11 @@
 import * as assert from 'assert';
 import * as path from 'path';
 import * as sinon from 'sinon';
-import { EnvironmentTypeComparer } from '../../client/interpreter/configuration/environmentTypeComparer';
+import {
+    EnvironmentTypeComparer,
+    EnvTypeHeuristic,
+    getEnvTypeHeuristic,
+} from '../../client/interpreter/configuration/environmentTypeComparer';
 import { IInterpreterHelper } from '../../client/interpreter/contracts';
 import { EnvironmentType, PythonEnvironment } from '../../client/pythonEnvironments/info';
 
@@ -225,3 +229,64 @@ suite('Environment sorting', () => {
         });
     });
 });
+
+suite('getEnvTypeHeuristic tests', () => {
+    const workspacePath = path.join('path', 'to', 'workspace');
+
+    const localGlobalEnvTypes = [
+        EnvironmentType.Venv,
+        EnvironmentType.Conda,
+        EnvironmentType.VirtualEnv,
+        EnvironmentType.VirtualEnvWrapper,
+        EnvironmentType.Pipenv,
+        EnvironmentType.Poetry,
+    ];
+
+    localGlobalEnvTypes.forEach((envType) => {
+        test('If the path to an environment starts with the workspace path it should be marked as local', () => {
+            const environment = {
+                envType,
+                envPath: path.join(workspacePath, 'my-environment'),
+                version: { major: 3, minor: 10, patch: 2 },
+            } as PythonEnvironment;
+
+            const envTypeHeuristic = getEnvTypeHeuristic(environment, workspacePath);
+
+            assert.strictEqual(envTypeHeuristic, EnvTypeHeuristic.Local);
+        });
+
+        test('If the path to an environment does not start with the workspace path it should be marked as global', () => {
+            const environment = {
+                envType,
+                envPath: path.join('path', 'to', 'my-environment'),
+                version: { major: 3, minor: 10, patch: 2 },
+            } as PythonEnvironment;
+
+            const envTypeHeuristic = getEnvTypeHeuristic(environment, workspacePath);
+
+            assert.strictEqual(envTypeHeuristic, EnvTypeHeuristic.Global);
+        });
+    });
+
+    const globalInterpretersEnvTypes = [
+        EnvironmentType.System,
+        EnvironmentType.WindowsStore,
+        EnvironmentType.Global,
+        EnvironmentType.Unknown,
+        EnvironmentType.Pyenv,
+    ];
+
+    globalInterpretersEnvTypes.forEach((envType) => {
+        test(`If the environment type is ${envType} and the environment path does not start with the workspace path it should be marked as a global interpreter`, () => {
+            const environment = {
+                envType,
+                envPath: path.join('path', 'to', 'a', 'global', 'interpreter'),
+                version: { major: 3, minor: 10, patch: 2 },
+            } as PythonEnvironment;
+
+            const envTypeHeuristic = getEnvTypeHeuristic(environment, workspacePath);
+
+            assert.strictEqual(envTypeHeuristic, EnvTypeHeuristic.GlobalInterpreters);
+        });
+    });
+});