Support for Create Env command to re-create env for venv (#21829)

Closes #21827

Author: karthiknadig

src/client/common/utils/localize.ts

@@ -464,6 +464,13 @@ export namespace CreateEnv {
         export const error = l10n.t('Creating virtual environment failed with error.');
         export const tomlExtrasQuickPickTitle = l10n.t('Select optional dependencies to install from pyproject.toml');
         export const requirementsQuickPickTitle = l10n.t('Select dependencies to install');
+        export const recreate = l10n.t('Recreate');
+        export const recreateDescription = l10n.t('Delete existing ".venv" environment and create a new one');
+        export const useExisting = l10n.t('Use Existing');
+        export const useExistingDescription = l10n.t('Use existing ".venv" environment with no changes to it');
+        export const existingVenvQuickPickPlaceholder = l10n.t('Use or Recreate existing environment?');
+        export const deletingEnvironmentProgress = l10n.t('Deleting existing ".venv" environment...');
+        export const errorDeletingEnvironment = l10n.t('Error while deleting existing ".venv" environment.');
     }
 
     export namespace Conda {

src/client/pythonEnvironments/creation/common/commonUtils.ts

@@ -1,9 +1,13 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+import * as fs from 'fs-extra';
+import * as path from 'path';
+import { WorkspaceFolder } from 'vscode';
 import { Commands } from '../../../common/constants';
 import { Common } from '../../../common/utils/localize';
 import { executeCommand } from '../../../common/vscodeApis/commandApis';
 import { showErrorMessage } from '../../../common/vscodeApis/windowApis';
+import { isWindows } from '../../../common/platform/platformService';
 
 export async function showErrorMessageWithLogs(message: string): Promise<void> {
     const result = await showErrorMessage(message, Common.openOutputPanel, Common.selectPythonInterpreter);
@@ -13,3 +17,18 @@ export async function showErrorMessageWithLogs(message: string): Promise<void> {
         await executeCommand(Commands.Set_Interpreter);
     }
 }
+
+export function getVenvPath(workspaceFolder: WorkspaceFolder): string {
+    return path.join(workspaceFolder.uri.fsPath, '.venv');
+}
+
+export async function hasVenv(workspaceFolder: WorkspaceFolder): Promise<boolean> {
+    return fs.pathExists(getVenvPath(workspaceFolder));
+}
+
+export function getVenvExecutable(workspaceFolder: WorkspaceFolder): string {
+    if (isWindows()) {
+        return path.join(getVenvPath(workspaceFolder), 'Scripts', 'python.exe');
+    }
+    return path.join(getVenvPath(workspaceFolder), 'bin', 'python');
+}

src/client/pythonEnvironments/creation/provider/venvCreationProvider.ts

@@ -17,8 +17,14 @@ import { MultiStepAction, MultiStepNode, withProgress } from '../../../common/vs
 import { sendTelemetryEvent } from '../../../telemetry';
 import { EventName } from '../../../telemetry/constants';
 import { VenvProgressAndTelemetry, VENV_CREATED_MARKER, VENV_EXISTING_MARKER } from './venvProgressAndTelemetry';
-import { showErrorMessageWithLogs } from '../common/commonUtils';
-import { IPackageInstallSelection, pickPackagesToInstall } from './venvUtils';
+import { getVenvExecutable, showErrorMessageWithLogs } from '../common/commonUtils';
+import {
+    ExistingVenvAction,
+    IPackageInstallSelection,
+    deleteEnvironment,
+    pickExistingVenvAction,
+    pickPackagesToInstall,
+} from './venvUtils';
 import { InputFlowAction } from '../../../common/utils/multiStepInput';
 import {
     CreateEnvironmentProvider,
@@ -150,33 +156,66 @@ export class VenvCreationProvider implements CreateEnvironmentProvider {
             undefined,
         );
 
-        let interpreter: string | undefined;
-        const interpreterStep = new MultiStepNode(
+        let existingVenvAction: ExistingVenvAction | undefined;
+        const existingEnvStep = new MultiStepNode(
             workspaceStep,
-            async () => {
-                if (workspace) {
+            async (context?: MultiStepAction) => {
+                if (workspace && context === MultiStepAction.Continue) {
                     try {
-                        interpreter = await this.interpreterQuickPick.getInterpreterViaQuickPick(
-                            workspace.uri,
-                            (i: PythonEnvironment) =>
-                                [
-                                    EnvironmentType.System,
-                                    EnvironmentType.MicrosoftStore,
-                                    EnvironmentType.Global,
-                                    EnvironmentType.Pyenv,
-                                ].includes(i.envType) && i.type === undefined, // only global intepreters
-                            {
-                                skipRecommended: true,
-                                showBackButton: true,
-                                placeholder: CreateEnv.Venv.selectPythonPlaceHolder,
-                                title: null,
-                            },
-                        );
+                        existingVenvAction = await pickExistingVenvAction(workspace);
+                        return MultiStepAction.Continue;
                     } catch (ex) {
-                        if (ex === InputFlowAction.back) {
+                        if (ex === MultiStepAction.Back || ex === MultiStepAction.Cancel) {
+                            return ex;
+                        }
+                        throw ex;
+                    }
+                } else if (context === MultiStepAction.Back) {
+                    return MultiStepAction.Back;
+                }
+                return MultiStepAction.Continue;
+            },
+            undefined,
+        );
+        workspaceStep.next = existingEnvStep;
+
+        let interpreter: string | undefined;
+        const interpreterStep = new MultiStepNode(
+            existingEnvStep,
+            async (context?: MultiStepAction) => {
+                if (workspace) {
+                    if (
+                        existingVenvAction === ExistingVenvAction.Recreate ||
+                        existingVenvAction === ExistingVenvAction.Create
+                    ) {
+                        try {
+                            interpreter = await this.interpreterQuickPick.getInterpreterViaQuickPick(
+                                workspace.uri,
+                                (i: PythonEnvironment) =>
+                                    [
+                                        EnvironmentType.System,
+                                        EnvironmentType.MicrosoftStore,
+                                        EnvironmentType.Global,
+                                        EnvironmentType.Pyenv,
+                                    ].includes(i.envType) && i.type === undefined, // only global intepreters
+                                {
+                                    skipRecommended: true,
+                                    showBackButton: true,
+                                    placeholder: CreateEnv.Venv.selectPythonPlaceHolder,
+                                    title: null,
+                                },
+                            );
+                        } catch (ex) {
+                            if (ex === InputFlowAction.back) {
+                                return MultiStepAction.Back;
+                            }
+                            interpreter = undefined;
+                        }
+                    } else if (existingVenvAction === ExistingVenvAction.UseExisting) {
+                        if (context === MultiStepAction.Back) {
                             return MultiStepAction.Back;
                         }
-                        interpreter = undefined;
+                        interpreter = getVenvExecutable(workspace);
                     }
                 }
 
@@ -189,7 +228,7 @@ export class VenvCreationProvider implements CreateEnvironmentProvider {
             },
             undefined,
         );
-        workspaceStep.next = interpreterStep;
+        existingEnvStep.next = interpreterStep;
 
         let addGitIgnore = true;
         let installPackages = true;
@@ -200,19 +239,23 @@ export class VenvCreationProvider implements CreateEnvironmentProvider {
         let installInfo: IPackageInstallSelection[] | undefined;
         const packagesStep = new MultiStepNode(
             interpreterStep,
-            async () => {
+            async (context?: MultiStepAction) => {
                 if (workspace && installPackages) {
-                    try {
-                        installInfo = await pickPackagesToInstall(workspace);
-                    } catch (ex) {
-                        if (ex === MultiStepAction.Back || ex === MultiStepAction.Cancel) {
-                            return ex;
+                    if (existingVenvAction !== ExistingVenvAction.UseExisting) {
+                        try {
+                            installInfo = await pickPackagesToInstall(workspace);
+                        } catch (ex) {
+                            if (ex === MultiStepAction.Back || ex === MultiStepAction.Cancel) {
+                                return ex;
+                            }
+                            throw ex;
                         }
-                        throw ex;
-                    }
-                    if (!installInfo) {
-                        traceVerbose('Virtual env creation exited during dependencies selection.');
-                        return MultiStepAction.Cancel;
+                        if (!installInfo) {
+                            traceVerbose('Virtual env creation exited during dependencies selection.');
+                            return MultiStepAction.Cancel;
+                        }
+                    } else if (context === MultiStepAction.Back) {
+                        return MultiStepAction.Back;
                     }
                 }
 
@@ -227,6 +270,32 @@ export class VenvCreationProvider implements CreateEnvironmentProvider {
             throw action;
         }
 
+        if (workspace) {
+            if (existingVenvAction === ExistingVenvAction.Recreate) {
+                sendTelemetryEvent(EventName.ENVIRONMENT_DELETE, undefined, {
+                    environmentType: 'venv',
+                    status: 'triggered',
+                });
+                if (await deleteEnvironment(workspace, interpreter)) {
+                    sendTelemetryEvent(EventName.ENVIRONMENT_DELETE, undefined, {
+                        environmentType: 'venv',
+                        status: 'deleted',
+                    });
+                } else {
+                    sendTelemetryEvent(EventName.ENVIRONMENT_DELETE, undefined, {
+                        environmentType: 'venv',
+                        status: 'failed',
+                    });
+                    throw MultiStepAction.Cancel;
+                }
+            } else if (existingVenvAction === ExistingVenvAction.UseExisting) {
+                sendTelemetryEvent(EventName.ENVIRONMENT_REUSE, undefined, {
+                    environmentType: 'venv',
+                });
+                return { path: getVenvExecutable(workspace), workspaceFolder: workspace };
+            }
+        }
+
         const args = generateCommandArgs(installInfo, addGitIgnore);
 
         return withProgress(

src/client/pythonEnvironments/creation/provider/venvDeleteUtils.ts

@@ -0,0 +1,99 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import * as fs from 'fs-extra';
+import * as path from 'path';
+import { WorkspaceFolder } from 'vscode';
+import { traceError, traceInfo } from '../../../logging';
+import { getVenvPath, showErrorMessageWithLogs } from '../common/commonUtils';
+import { CreateEnv } from '../../../common/utils/localize';
+import { sleep } from '../../../common/utils/async';
+import { switchSelectedPython } from './venvSwitchPython';
+
+async function tryDeleteFile(file: string): Promise<boolean> {
+    try {
+        if (!(await fs.pathExists(file))) {
+            return true;
+        }
+        await fs.unlink(file);
+        return true;
+    } catch (err) {
+        traceError(`Failed to delete file [${file}]:`, err);
+        return false;
+    }
+}
+
+async function tryDeleteDir(dir: string): Promise<boolean> {
+    try {
+        if (!(await fs.pathExists(dir))) {
+            return true;
+        }
+        await fs.rmdir(dir, {
+            recursive: true,
+            maxRetries: 10,
+            retryDelay: 200,
+        });
+        return true;
+    } catch (err) {
+        traceError(`Failed to delete directory [${dir}]:`, err);
+        return false;
+    }
+}
+
+export async function deleteEnvironmentNonWindows(workspaceFolder: WorkspaceFolder): Promise<boolean> {
+    const venvPath = getVenvPath(workspaceFolder);
+    if (await tryDeleteDir(venvPath)) {
+        traceInfo(`Deleted venv dir: ${venvPath}`);
+        return true;
+    }
+    showErrorMessageWithLogs(CreateEnv.Venv.errorDeletingEnvironment);
+    return false;
+}
+
+export async function deleteEnvironmentWindows(
+    workspaceFolder: WorkspaceFolder,
+    interpreter: string | undefined,
+): Promise<boolean> {
+    const venvPath = getVenvPath(workspaceFolder);
+    const venvPythonPath = path.join(venvPath, 'Scripts', 'python.exe');
+
+    if (await tryDeleteFile(venvPythonPath)) {
+        traceInfo(`Deleted python executable: ${venvPythonPath}`);
+        if (await tryDeleteDir(venvPath)) {
+            traceInfo(`Deleted ".venv" dir: ${venvPath}`);
+            return true;
+        }
+
+        traceError(`Failed to delete ".venv" dir: ${venvPath}`);
+        traceError(
+            'This happens if the virtual environment is still in use, or some binary in the venv is still running.',
+        );
+        traceError(`Please delete the ".venv" manually: [${venvPath}]`);
+        showErrorMessageWithLogs(CreateEnv.Venv.errorDeletingEnvironment);
+        return false;
+    }
+    traceError(`Failed to delete python executable: ${venvPythonPath}`);
+    traceError('This happens if the virtual environment is still in use.');
+
+    if (interpreter) {
+        traceError('We will attempt to switch python temporarily to delete the ".venv"');
+
+        await switchSelectedPython(interpreter, workspaceFolder.uri, 'temporarily to delete the ".venv"');
+
+        traceInfo(`Attempting to delete ".venv" again: ${venvPath}`);
+        const ms = 500;
+        for (let i = 0; i < 5; i = i + 1) {
+            traceInfo(`Waiting for ${ms}ms to let processes exit, before a delete attempt.`);
+            await sleep(ms);
+            if (await tryDeleteDir(venvPath)) {
+                traceInfo(`Deleted ".venv" dir: ${venvPath}`);
+                return true;
+            }
+            traceError(`Failed to delete ".venv" dir [${venvPath}] (attempt ${i + 1}/5).`);
+        }
+    } else {
+        traceError(`Please delete the ".venv" dir manually: [${venvPath}]`);
+    }
+    showErrorMessageWithLogs(CreateEnv.Venv.errorDeletingEnvironment);
+    return false;
+}

src/client/pythonEnvironments/creation/provider/venvSwitchPython.ts

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import * as path from 'path';
+import { Disposable, Uri } from 'vscode';
+import { createDeferred } from '../../../common/utils/async';
+import { getExtension } from '../../../common/vscodeApis/extensionsApi';
+import { PVSC_EXTENSION_ID, PythonExtension } from '../../../api/types';
+import { traceInfo } from '../../../logging';
+
+export async function switchSelectedPython(interpreter: string, uri: Uri, purpose: string): Promise<void> {
+    let dispose: Disposable | undefined;
+    try {
+        const deferred = createDeferred<void>();
+        const api: PythonExtension = getExtension(PVSC_EXTENSION_ID)?.exports as PythonExtension;
+        dispose = api.environments.onDidChangeActiveEnvironmentPath(async (e) => {
+            if (path.normalize(e.path) === path.normalize(interpreter)) {
+                traceInfo(`Switched to interpreter ${purpose}: ${interpreter}`);
+                deferred.resolve();
+            }
+        });
+        api.environments.updateActiveEnvironmentPath(interpreter, uri);
+        traceInfo(`Switching interpreter ${purpose}: ${interpreter}`);
+        await deferred.promise;
+    } finally {
+        dispose?.dispose();
+    }
+}