Define preferred Python Env in new Kernel Picker (#14051)

src/notebooks/controllers/kernelSource/localPythonKernelSelector.node.ts

@@ -13,27 +13,34 @@ import {
     pythonEnvironmentQuickPick
 } from '../../../platform/interpreter/pythonEnvironmentPicker.node';
 import { BaseProviderBasedQuickPick } from '../../../platform/common/providerBasedQuickPick';
-import { Environment } from '../../../platform/api/pythonApiTypes';
+import { Environment, ProposedExtensionAPI } from '../../../platform/api/pythonApiTypes';
 import { DataScience } from '../../../platform/common/utils/localize';
 import { PythonEnvKernelConnectionCreator } from '../pythonEnvKernelConnectionCreator.node';
 import { IPythonApiProvider, IPythonExtensionChecker } from '../../../platform/api/types';
 import { PythonEnvironmentQuickPickItemProvider } from '../../../platform/interpreter/pythonEnvironmentQuickPickProvider.node';
 import { Disposables } from '../../../platform/common/utils';
 import { PythonEnvironmentFilter } from '../../../platform/interpreter/filter/filterService';
+import { noop } from '../../../platform/common/utils/misc';
+import { findPreferredPythonEnvironment } from '../preferredKernelConnectionService.node';
 
 export class LocalPythonKernelSelector extends Disposables {
     private readonly pythonEnvPicker: BaseProviderBasedQuickPick<Environment>;
+    private readonly provider: PythonEnvironmentQuickPickItemProvider;
+    private pythonApi?: ProposedExtensionAPI;
     constructor(
         private readonly notebook: NotebookDocument,
         private readonly token: CancellationToken
     ) {
         super();
         const filter = ServiceContainer.instance.get<PythonEnvironmentFilter>(PythonEnvironmentFilter);
-        const provider = ServiceContainer.instance
+        const pythonExtensionChecker = ServiceContainer.instance.get<IPythonExtensionChecker>(IPythonExtensionChecker);
+        const pythonApiProvider = ServiceContainer.instance.get<IPythonApiProvider>(IPythonApiProvider);
+
+        this.provider = ServiceContainer.instance
             .get<PythonEnvironmentQuickPickItemProvider>(PythonEnvironmentQuickPickItemProvider)
             .withFilter((item) => !filter.isPythonEnvironmentExcluded(item));
         this.pythonEnvPicker = new BaseProviderBasedQuickPick(
-            provider,
+            this.provider,
             pythonEnvironmentQuickPick,
             getPythonEnvironmentCategory,
             { supportsBack: true }
@@ -43,6 +50,33 @@ export class LocalPythonKernelSelector extends Disposables {
             { label: `$(add) ${DataScience.createPythonEnvironmentInQuickPick}` },
             this.createNewEnvironment.bind(this)
         );
+        const computePreferredEnv = () => {
+            if (!this.pythonApi || token.isCancellationRequested) {
+                return;
+            }
+            this.pythonEnvPicker.recommended = findPreferredPythonEnvironment(this.notebook, this.pythonApi);
+            console.log(1234);
+        };
+        const setupApi = (api?: ProposedExtensionAPI) => {
+            if (!api) {
+                return;
+            }
+            this.pythonApi = api;
+            computePreferredEnv();
+            this.disposables.push(api.environments.onDidChangeActiveEnvironmentPath(computePreferredEnv));
+            this.disposables.push(api.environments.onDidChangeEnvironments(computePreferredEnv));
+        };
+        if (pythonExtensionChecker.isPythonExtensionInstalled) {
+            pythonApiProvider.getNewApi().then(setupApi).catch(noop);
+        } else {
+            pythonExtensionChecker.onPythonExtensionInstallationStatusChanged(
+                () => pythonApiProvider.getNewApi().then(setupApi),
+                this,
+                this.disposables
+            );
+        }
+
+        computePreferredEnv();
     }
 
     public async selectKernel(): Promise<

src/notebooks/controllers/preferredKernelConnectionService.node.ts

@@ -0,0 +1,68 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { NotebookDocument, workspace, Uri } from 'vscode';
+import * as path from '../../platform/vscode-path/resources';
+import { isParentPath } from '../../platform/common/platform/fileUtils';
+import { EnvironmentType } from '../../platform/pythonEnvironments/info';
+import { getEnvironmentType } from '../../platform/interpreter/helpers';
+import { Environment, ProposedExtensionAPI } from '../../platform/api/pythonApiTypes';
+
+export function findPreferredPythonEnvironment(
+    notebook: NotebookDocument,
+    pythonApi: ProposedExtensionAPI
+): Environment | undefined {
+    // 1. Check if we have a .conda or .venv virtual env in the local workspace folder.
+    const localEnv = findPythonEnvironmentClosestToNotebook(notebook, pythonApi.environments.known);
+    if (localEnv) {
+        return localEnv;
+    }
+
+    const findMatchingActiveEnvironment = () => {
+        const envPath = pythonApi.environments.getActiveEnvironmentPath(notebook.uri);
+        return envPath && pythonApi.environments.known.find((e) => e.id === envPath.id);
+    };
+    // 3. Fall back to the active interpreter.
+    return findMatchingActiveEnvironment();
+}
+
+function findPythonEnvironmentClosestToNotebook(notebook: NotebookDocument, envs: readonly Environment[]) {
+    const defaultFolder =
+        workspace.getWorkspaceFolder(notebook.uri)?.uri ||
+        (workspace.workspaceFolders?.length === 1 ? workspace.workspaceFolders[0].uri : undefined);
+    const localEnvNextToNbFile = findPythonEnvBelongingToFolder(path.dirname(notebook.uri), envs);
+    if (localEnvNextToNbFile) {
+        return localEnvNextToNbFile;
+    }
+    if (defaultFolder) {
+        return findPythonEnvBelongingToFolder(defaultFolder, envs);
+    }
+}
+
+function findPythonEnvBelongingToFolder(folder: Uri, pythonEnvs: readonly Environment[]) {
+    const localEnvs = pythonEnvs.filter((p) =>
+        // eslint-disable-next-line local-rules/dont-use-fspath
+        isParentPath(p.environment?.folderUri?.fsPath || p.executable.uri?.fsPath || p.path, folder.fsPath)
+    );
+
+    // Find an environment that is a .venv or .conda environment.
+    // Give preference to .venv over .conda.
+    // & give preference to .venv or .conda over any other environment.
+    return localEnvs.find(
+        (e) => getEnvironmentType(e) === EnvironmentType.Venv && e.environment?.name?.toLowerCase() === '.venv'
+    ) ||
+        localEnvs.find(
+            (e) => getEnvironmentType(e) === EnvironmentType.Conda && e.environment?.name?.toLowerCase() === '.conda'
+        ) ||
+        localEnvs.find(
+            (e) =>
+                [EnvironmentType.VirtualEnv, EnvironmentType.VirtualEnvWrapper].includes(getEnvironmentType(e)) &&
+                e.environment?.name?.toLowerCase() === '.venv'
+        ) ||
+        localEnvs.find(
+            (e) => e.environment?.name?.toLowerCase() === '.venv' || e.environment?.name?.toLowerCase() === '.conda'
+        ) ||
+        localEnvs.length
+        ? localEnvs[0]
+        : undefined;
+}

src/platform/api/pythonApiTypes.ts

@@ -145,7 +145,7 @@ export type Environment = EnvironmentPath & {
     /**
      * Carries Python version information known at this moment.
      */
-    readonly version: VersionInfo & {
+    readonly version?: VersionInfo & {
         /**
          * Value of `sys.version` in sys module if known at this moment.
          */

src/platform/common/providerBasedQuickPick.ts

@@ -83,7 +83,11 @@ export class BaseProviderBasedQuickPick<T extends { id: string }> extends Dispos
     private commands = new Set<CommandQuickPickItem<T>>();
     private _recommended?: T;
     public set recommended(item: T | undefined) {
+        const changed = this._recommended?.id !== item?.id;
         this._recommended = item;
+        if (changed && this.quickPick) {
+            this.rebuildQuickPickItems(this.quickPick);
+        }
     }
     public get recommended() {
         return this._recommended;

src/platform/interpreter/helpers.ts

@@ -8,9 +8,17 @@ import { Environment } from '../api/pythonApiTypes';
 
 export function getPythonEnvDisplayName(interpreter: PythonEnvironment | Environment) {
     if ('executable' in interpreter) {
-        const version = interpreter.version?.major
-            ? `${interpreter.version.major}.${interpreter.version.minor}.${interpreter.version.micro}`
-            : '';
+        const versionParts: string[] = [];
+        if (typeof interpreter.version?.major === 'number') {
+            versionParts.push(interpreter.version.major.toString());
+            if (typeof interpreter.version.minor === 'number') {
+                versionParts.push(interpreter.version.minor.toString());
+                if (typeof interpreter.version.micro === 'number') {
+                    versionParts.push(interpreter.version.micro.toString());
+                }
+            }
+        }
+        const version = versionParts.length ? versionParts.join('.') : '';
         const envName = interpreter.environment ? basename(interpreter.environment?.folderUri) : '';
         const nameWithVersion = version ? `Python ${version}` : 'Python';
         if (isCondaEnvironmentWithoutPython(interpreter) && envName) {

src/platform/interpreter/pythonEnvironmentPicker.node.ts

@@ -14,7 +14,7 @@ export function pythonEnvironmentQuickPick(item: Environment, quickPick: BasePro
     const label = getPythonEnvDisplayName(item);
     const icon =
         item.id === quickPick.recommended?.id
-            ? ' $(star-full) '
+            ? '$(star-full) '
             : isCondaEnvironmentWithoutPython(item)
             ? '$(warning) '
             : '';