Configure Python Env tool

Author: DonJayamanne

src/client/chat/configurePythonEnvTool.ts

@@ -43,7 +43,7 @@ export function registerTools(
         ),
     );
     ourTools.add(
-        lm.registerTool(ConfigurePythonEnvTool.toolName, new ConfigurePythonEnvTool(environmentsApi, discoverApi, serviceContainer)),
+        lm.registerTool(ConfigurePythonEnvTool.toolName, new ConfigurePythonEnvTool(environmentsApi, serviceContainer)),
     );
     ourTools.add(
         extensions.onDidChange(() => {

src/client/chat/index.ts

@@ -45,7 +45,7 @@ import { DebugService } from './common/application/debugService';
 import { DebugSessionEventDispatcher } from './debugger/extension/hooks/eventHandlerDispatcher';
 import { IDebugSessionEventHandlers } from './debugger/extension/hooks/types';
 import { WorkspaceService } from './common/application/workspace';
-import { IInterpreterQuickPick } from './interpreter/configuration/types';
+import { IInterpreterQuickPick, IRecommendedEnvironmentService } from './interpreter/configuration/types';
 import { registerAllCreateEnvironmentFeatures } from './pythonEnvironments/creation/registrations';
 import { registerCreateEnvironmentTriggers } from './pythonEnvironments/creation/createEnvironmentTrigger';
 import { initializePersistentStateForTriggers } from './common/persistentState';
@@ -111,6 +111,7 @@ export function activateFeatures(ext: ExtensionState, _components: Components):
         interpreterPathService,
         interpreterService,
         pathUtils,
+        ext.legacyIOC.serviceContainer.get<IRecommendedEnvironmentService>(IRecommendedEnvironmentService),
     );
     const executionHelper = ext.legacyIOC.serviceContainer.get<ICodeExecutionHelper>(ICodeExecutionHelper);
     const commandManager = ext.legacyIOC.serviceContainer.get<ICommandManager>(ICommandManager);

src/client/extensionActivation.ts

@@ -569,7 +569,10 @@ export class SetInterpreterCommand extends BaseInterpreterSelectorCommand implem
      * @returns true when an interpreter was set, undefined if the user cancelled the quickpick.
      */
     @captureTelemetry(EventName.SELECT_INTERPRETER)
-    public async setInterpreter(): Promise<true | undefined> {
+    public async setInterpreter(options?: {
+        hideCreateVenv?: boolean;
+        showBackButton?: boolean;
+    }): Promise<SelectEnvironmentResult | undefined> {
         const targetConfig = await this.getConfigTargets();
         if (!targetConfig) {
             return;
@@ -578,11 +581,25 @@ export class SetInterpreterCommand extends BaseInterpreterSelectorCommand implem
         const wkspace = targetConfig[0].folderUri;
         const interpreterState: InterpreterStateArgs = { path: undefined, workspace: wkspace };
         const multiStep = this.multiStepFactory.create<InterpreterStateArgs>();
-        await multiStep.run(
-            (input, s) => this._pickInterpreter(input, s, undefined, { showCreateEnvironment: true }),
-            interpreterState,
-        );
-
+        try {
+            await multiStep.run(
+                (input, s) =>
+                    this._pickInterpreter(input, s, undefined, {
+                        showCreateEnvironment: !options?.hideCreateVenv,
+                        showBackButton: options?.showBackButton,
+                    }),
+                interpreterState,
+            );
+        } catch (ex) {
+            if (ex === InputFlowAction.back) {
+                // User clicked back button, so we need to return this action.
+                return { action: 'Back' };
+            }
+            if (ex === InputFlowAction.cancel) {
+                // User clicked cancel button, so we need to return this action.
+                return { action: 'Cancel' };
+            }
+        }
         if (interpreterState.path !== undefined) {
             // User may choose to have an empty string stored, so variable `interpreterState.path` may be
             // an empty string, in which case we should update.
@@ -591,7 +608,7 @@ export class SetInterpreterCommand extends BaseInterpreterSelectorCommand implem
             if (useEnvExtension()) {
                 await setInterpreterLegacy(interpreterState.path, wkspace);
             }
-            return true;
+            return { path: interpreterState.path };
         }
     }
 
@@ -692,3 +709,14 @@ function getGroup(item: IInterpreterQuickPickItem, workspacePath?: string) {
             return EnvGroups[item.interpreter.envType];
     }
 }
+
+export type SelectEnvironmentResult = {
+    /**
+     * Path to the executable python in the environment
+     */
+    readonly path?: string;
+    /*
+     * User action that resulted in exit from the create environment flow.
+     */
+    readonly action?: 'Back' | 'Cancel';
+};

src/client/interpreter/configuration/interpreterSelector/commands/setInterpreter.ts

@@ -3,7 +3,7 @@
 
 import { inject, injectable } from 'inversify';
 import { IRecommendedEnvironmentService } from './types';
-import { PythonExtension } from '../../api/types';
+import { PythonExtension, ResolvedEnvironment } from '../../api/types';
 import { IExtensionContext, Resource } from '../../common/types';
 import { Uri, workspace } from 'vscode';
 import { getWorkspaceStateValue, updateWorkspaceStateValue } from '../../common/persistentState';
@@ -32,7 +32,38 @@ export class RecommendedEnvironmentService implements IRecommendedEnvironmentSer
         }
     }
 
-    getRecommededEnvironment(
+    async getRecommededEnvironment(
+        resource: Resource,
+    ): Promise<
+        | {
+              environment: ResolvedEnvironment;
+              reason: 'globalUserSelected' | 'workspaceUserSelected' | 'defaultRecommended';
+          }
+        | undefined
+    > {
+        if (!workspace.isTrusted || !this.api) {
+            return undefined;
+        }
+        const preferred = await this.getRecommededInternal(resource);
+        if (!preferred) {
+            return undefined;
+        }
+        const activeEnv = await this.api.resolveEnvironment(this.api.getActiveEnvironmentPath(resource));
+        const recommendedEnv = await this.api.resolveEnvironment(preferred.environmentPath);
+        if (activeEnv && recommendedEnv && activeEnv.id !== recommendedEnv.id) {
+            traceError(
+                `Active environment ${activeEnv.id} is different from recommended environment ${
+                    recommendedEnv.id
+                } for resource ${resource?.toString()}`,
+            );
+            return undefined;
+        }
+        if (recommendedEnv) {
+            return { environment: recommendedEnv, reason: preferred.reason };
+        }
+        return undefined;
+    }
+    getRecommededInternal(
         resource: Resource,
     ):
         | { environmentPath: string; reason: 'globalUserSelected' | 'workspaceUserSelected' | 'defaultRecommended' }

src/client/interpreter/configuration/recommededEnvironmentService.ts

@@ -1,7 +1,7 @@
 import { ConfigurationTarget, Disposable, QuickPickItem, Uri } from 'vscode';
 import { Resource } from '../../common/types';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
-import { PythonExtension } from '../../api/types';
+import { PythonExtension, ResolvedEnvironment } from '../../api/types';
 
 export interface IPythonPathUpdaterService {
     updatePythonPath(pythonPath: string | undefined): Promise<void>;
@@ -104,7 +104,11 @@ export interface IRecommendedEnvironmentService {
     trackUserSelectedEnvironment(environmentPath: string | undefined, uri: Uri | undefined): void;
     getRecommededEnvironment(
         resource: Resource,
-    ):
-        | { environmentPath: string; reason: 'globalUserSelected' | 'workspaceUserSelected' | 'defaultRecommended' }
-        | undefined;
+    ): Promise<
+        | {
+              environment: ResolvedEnvironment;
+              reason: 'globalUserSelected' | 'workspaceUserSelected' | 'defaultRecommended';
+          }
+        | undefined
+    >;
 }

src/client/interpreter/configuration/types.ts

@@ -72,7 +72,7 @@ type PythonApiForJupyterExtension = {
      * Returns the preferred environment for the given URI.
      */
     getRecommededEnvironment(
-        uri: Uri,
+        uri: Uri | undefined,
     ): Promise<
         | {
               environment: EnvironmentPath;
@@ -147,14 +147,7 @@ export class JupyterExtensionIntegration {
                 if (!this.environmentApi) {
                     return undefined;
                 }
-                const preferred = this.preferredEnvironmentService.getRecommededEnvironment(uri);
-                if (!preferred) {
-                    return undefined;
-                }
-                const environment = workspace.isTrusted
-                    ? await this.environmentApi.resolveEnvironment(preferred.environmentPath)
-                    : undefined;
-                return environment ? { environment, reason: preferred.reason } : undefined;
+                return this.preferredEnvironmentService.getRecommededEnvironment(uri);
             },
         });
         return undefined;

src/client/jupyter/jupyterIntegration.ts

@@ -5,7 +5,7 @@ import { ConfigurationTarget, Disposable, QuickInputButtons } from 'vscode';
 import { Commands } from '../../common/constants';
 import { IDisposableRegistry, IInterpreterPathService, IPathUtils } from '../../common/types';
 import { executeCommand, registerCommand } from '../../common/vscodeApis/commandApis';
-import { IInterpreterQuickPick } from '../../interpreter/configuration/types';
+import { IInterpreterQuickPick, IRecommendedEnvironmentService } from '../../interpreter/configuration/types';
 import { getCreationEvents, handleCreateEnvironmentCommand } from './createEnvironment';
 import { condaCreationProvider } from './provider/condaCreationProvider';
 import { VenvCreationProvider } from './provider/venvCreationProvider';
@@ -63,6 +63,7 @@ export function registerCreateEnvironmentFeatures(
     interpreterQuickPick: IInterpreterQuickPick,
     interpreterPathService: IInterpreterPathService,
     pathUtils: IPathUtils,
+    recommededEnvService: IRecommendedEnvironmentService,
 ): void {
     disposables.push(
         registerCommand(
@@ -108,6 +109,7 @@ export function registerCreateEnvironmentFeatures(
                     ConfigurationTarget.WorkspaceFolder,
                     e.path,
                 );
+                recommededEnvService.trackUserSelectedEnvironment(e.path, e.workspaceFolder?.uri);
                 showInformationMessage(`${CreateEnv.informEnvCreation} ${pathUtils.getDisplayName(e.path)}`);
             }
         }),

src/client/pythonEnvironments/creation/createEnvApi.ts

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 import { IDisposableRegistry, IInterpreterPathService, IPathUtils } from '../../common/types';
-import { IInterpreterQuickPick } from '../../interpreter/configuration/types';
+import { IInterpreterQuickPick, IRecommendedEnvironmentService } from '../../interpreter/configuration/types';
 import { IInterpreterService } from '../../interpreter/contracts';
 import { registerCreateEnvironmentFeatures } from './createEnvApi';
 import { registerCreateEnvironmentButtonFeatures } from './createEnvButtonContext';
@@ -16,8 +16,9 @@ export function registerAllCreateEnvironmentFeatures(
     interpreterPathService: IInterpreterPathService,
     interpreterService: IInterpreterService,
     pathUtils: IPathUtils,
+    recommededEnvService: IRecommendedEnvironmentService
 ): void {
-    registerCreateEnvironmentFeatures(disposables, interpreterQuickPick, interpreterPathService, pathUtils);
+    registerCreateEnvironmentFeatures(disposables, interpreterQuickPick, interpreterPathService, pathUtils, recommededEnvService);
     registerCreateEnvironmentButtonFeatures(disposables);
     registerPyProjectTomlFeatures(disposables);
     registerInstalledPackagesDiagnosticsProvider(disposables, interpreterService);

src/client/pythonEnvironments/creation/registrations.ts

@@ -14,6 +14,7 @@ import { registerCreateEnvironmentFeatures } from '../../../client/pythonEnviron
 import * as windowApis from '../../../client/common/vscodeApis/windowApis';
 import { handleCreateEnvironmentCommand } from '../../../client/pythonEnvironments/creation/createEnvironment';
 import { CreateEnvironmentProvider } from '../../../client/pythonEnvironments/creation/proposed.createEnvApis';
+import { noop } from '../../core';
 
 chaiUse(chaiAsPromised.default);
 
@@ -48,6 +49,11 @@ suite('Create Environment APIs', () => {
             interpreterQuickPick.object,
             interpreterPathService.object,
             pathUtils.object,
+            {
+                getRecommededEnvironment: () => Promise.resolve(undefined),
+                registerEnvApi: noop,
+                trackUserSelectedEnvironment: noop,
+            },
         );
     });
     teardown(() => {