Simplify UI for selecting virtual environments (#1413)

## Changes Create our own environment selector quick pick: - Explicit "Select Python Environment" title, not "Select Interpreter" that the python extension uses - Show only environments, not all interpreters as the python extension does - Still show the option to use python extension selector for power users New selector: <img width="952" alt="Screenshot 2024-10-29 at 10 53 21" src="https://github.com/user-attachments/assets/757e5566-c5fb-419a-8480-68c09632cfab"> Old selector, or when you click on the "Use Python Extension to setup environments": <img width="952" alt="Screenshot 2024-10-29 at 10 53 30" src="https://github.com/user-attachments/assets/feee5575-4c0c-486b-8f77-75e4d5108843"> As before, when there's no detected environments we show "create" dialog directly: <img width="952" alt="Screenshot 2024-10-29 at 11 06 04" src="https://github.com/user-attachments/assets/b87ad8f0-fc7f-4702-a557-8833d394f19f"> ## Tests Manually
databricks · Oct 29, 2024 · d1b8508 · d1b8508
1 parent faa901d
commit d1b8508
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 5 deletions.
diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
@@ -276,6 +276,13 @@
                 "enablement": "databricks.context.activated && databricks.context.loggedIn",
                 "category": "Databricks"
             },
+            {
+                "command": "databricks.environment.setup",
+                "title": "Setup python environment",
+                "icon": "$(gear)",
+                "enablement": "databricks.context.activated && databricks.context.loggedIn",
+                "category": "Databricks"
+            },
             {
                 "command": "databricks.environment.selectPythonInterpreter",
                 "title": "Change Python environment",

diff --git a/packages/databricks-vscode/src/language/EnvironmentCommands.ts b/packages/databricks-vscode/src/language/EnvironmentCommands.ts
@@ -1,8 +1,10 @@
-import {window, commands} from "vscode";
+import {window, commands, QuickPickItem} from "vscode";
 import {FeatureManager} from "../feature-manager/FeatureManager";
 import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
 import {Cluster} from "../sdk-extensions";
 import {EnvironmentDependenciesInstaller} from "./EnvironmentDependenciesInstaller";
+import {Environment} from "./MsPythonExtensionApi";
+import {environmentName} from "../utils/environmentUtils";
 
 export class EnvironmentCommands {
     constructor(
@@ -55,12 +57,44 @@ export class EnvironmentCommands {
         const environments =
             await this.pythonExtension.getAvailableEnvironments();
         if (environments.length > 0) {
-            await this.pythonExtension.selectPythonInterpreter();
+            await this.showEnvironmentsQuickPick(environments);
         } else {
             await this.pythonExtension.createPythonEnvironment();
         }
     }
 
+    async showEnvironmentsQuickPick(environments: Environment[]) {
+        const envPicks: (QuickPickItem & {path?: string})[] = environments.map(
+            (env) => ({
+                label: environmentName(env),
+                description: env.path,
+                path: env.path,
+            })
+        );
+        const createNewLabel = "$(add) Create new environment";
+        const usePythonExtensionLabel =
+            "$(gear) Use Python Extension to setup environments";
+        const staticPicks: QuickPickItem[] = [
+            {label: createNewLabel, alwaysShow: true},
+            {label: usePythonExtensionLabel, alwaysShow: true},
+        ];
+        const selectedPick = await window.showQuickPick(
+            envPicks.concat(staticPicks),
+            {title: "Select Python Environment"}
+        );
+        if (selectedPick) {
+            if (selectedPick.label === createNewLabel) {
+                await this.pythonExtension.createPythonEnvironment();
+            } else if (selectedPick.label === usePythonExtensionLabel) {
+                await this.pythonExtension.selectPythonInterpreter();
+            } else if (selectedPick.path) {
+                await this.pythonExtension.api.environments.updateActiveEnvironmentPath(
+                    selectedPick.path
+                );
+            }
+        }
+    }
+
     async reinstallDBConnect(cluster?: Cluster) {
         const state = await this.featureManager.isEnabled(
             "environment.dependencies"

diff --git a/packages/databricks-vscode/src/language/EnvironmentDependenciesInstaller.ts b/packages/databricks-vscode/src/language/EnvironmentDependenciesInstaller.ts
@@ -1,4 +1,4 @@
-import {EventEmitter, OutputChannel, window} from "vscode";
+import {commands, EventEmitter, OutputChannel, window} from "vscode";
 
 import {Disposable} from "vscode";
 import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
@@ -96,7 +96,10 @@ export class EnvironmentDependenciesInstaller implements Disposable {
             case "Change version":
                 return this.installWithVersionPrompt();
             case "Change environment":
-                return this.pythonExtension.selectPythonInterpreter();
+                await commands.executeCommand(
+                    "databricks.environment.selectPythonInterpreter"
+                );
+                return;
         }
     }
 }
diff --git a/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts b/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts
@@ -60,7 +60,15 @@ export class MsPythonExtensionWrapper implements Disposable {
 
     async getAvailableEnvironments() {
         await this.api.environments.refreshEnvironments();
-        return this.api.environments.known.filter((env) => env.environment);
+        const filteredEnvs = [];
+        for (const env of this.api.environments.known) {
+            const resolvedEnv =
+                await this.api.environments.resolveEnvironment(env);
+            if (resolvedEnv && resolvedEnv.environment) {
+                filteredEnvs.push(env);
+            }
+        }
+        return filteredEnvs;
     }
 
     get onDidChangePythonExecutable(): Event<Uri | undefined> {

diff --git a/packages/databricks-vscode/src/utils/environmentUtils.ts b/packages/databricks-vscode/src/utils/environmentUtils.ts
@@ -0,0 +1,9 @@
+import {Environment} from "../language/MsPythonExtensionApi";
+
+export function environmentName(env: Environment) {
+    const version = env.version
+        ? `${env.version.major}.${env.version.minor}.${env.version.micro} `
+        : "";
+    const name = env.environment?.name ?? env.path;
+    return `${version}${name}`;
+}