fix: add load check to more places

Author: karthiknadig

src/api.ts

@@ -333,7 +333,7 @@ export interface QuickCreateConfig {
  */
 export interface EnvironmentManager {
     /**
-     * The name of the environment manager.
+     * The name of the environment manager. Allowed characters (a-z, A-Z, 0-9, -, _).
      */
     readonly name: string;
 
@@ -564,7 +564,7 @@ export interface DidChangePackagesEventArgs {
  */
 export interface PackageManager {
     /**
-     * The name of the package manager.
+     * The name of the package manager. Allowed characters (a-z, A-Z, 0-9, -, _).
      */
     name: string;
 

src/common/localize.ts

@@ -14,6 +14,10 @@ export namespace Common {
     export const quickCreate = l10n.t('Quick Create');
 }
 
+export namespace WorkbenchStrings {
+    export const installExtension = l10n.t('Install Extension');
+}
+
 export namespace Interpreter {
     export const statusBarSelect = l10n.t('Select Interpreter');
     export const browsePath = l10n.t('Browse...');

src/common/workbenchCommands.ts

@@ -0,0 +1,12 @@
+import { commands, Uri } from 'vscode';
+
+export async function installExtension(
+    extensionId: Uri | string,
+    options?: {
+        installOnlyNewlyAddedFromExtensionPackVSIX?: boolean;
+        installPreReleaseVersion?: boolean;
+        donotSync?: boolean;
+    },
+): Promise<void> {
+    await commands.executeCommand('workbench.extensions.installExtension', extensionId, options);
+}

src/features/common/managerReady.ts

@@ -1,14 +1,26 @@
-import { Disposable, Uri } from 'vscode';
+import { Disposable, l10n, Uri } from 'vscode';
 import { EnvironmentManagers, PythonProjectManager } from '../../internal.api';
 import { createDeferred, Deferred } from '../../common/utils/deferred';
 import { allExtensions } from '../../common/extension.apis';
-import { traceError } from '../../common/logging';
+import { traceError, traceInfo } from '../../common/logging';
 import { showErrorMessage } from '../../common/window.apis';
 import { getDefaultEnvManagerSetting, getDefaultPkgManagerSetting } from '../settings/settingHelpers';
+import { WorkbenchStrings } from '../../common/localize';
+import { installExtension } from '../../common/workbenchCommands';
 
 interface ManagerReady extends Disposable {
     waitForEnvManager(uris?: Uri[]): Promise<void>;
+    waitForEnvManagerId(managerIds: string[]): Promise<void>;
+    waitForAllEnvManagers(): Promise<void>;
     waitForPkgManager(uris?: Uri[]): Promise<void>;
+    waitForPkgManagerId(managerIds: string[]): Promise<void>;
+}
+
+function getExtensionId(managerId: string): string | undefined {
+    // format <extension-id>:<manager-name>
+    const regex = /^(.*):([a-zA-Z0-9-_]*)$/;
+    const parts = regex.exec(managerId);
+    return parts ? parts[1] : undefined;
 }
 
 class ManagerReadyImpl implements ManagerReady {
@@ -44,10 +56,30 @@ class ManagerReadyImpl implements ManagerReady {
         const installed = allExtensions().some((ext) => managerId.startsWith(`${ext.id}:`));
         if (!installed && !this.checked.has(managerId)) {
             this.checked.add(managerId);
-            traceError(`Extension for manager ${managerId} is not installed.`);
-            showErrorMessage(`Extension for manager ${managerId} is not installed.`);
+            const extId = getExtensionId(managerId);
+            setImmediate(async () => {
+                if (extId) {
+                    traceError(`Extension for manager ${extId} is not installed.`);
+                    const result = await showErrorMessage(
+                        l10n.t(`Extension for {0} is not installed or enabled for this workspace.`, extId),
+                        WorkbenchStrings.installExtension,
+                    );
+                    if (result === WorkbenchStrings.installExtension) {
+                        traceInfo(`Installing extension: ${extId}`);
+                        try {
+                            await installExtension(extId);
+                            traceInfo(`Extension ${extId} installed.`);
+                        } catch (err) {
+                            traceError(`Failed to install extension: ${extId}`, err);
+                        }
+                    }
+                } else {
+                    showErrorMessage(
+                        l10n.t(`Extension for {0} is not installed or enabled for this workspace.`, managerId),
+                    );
+                }
+            });
         }
-        return installed;
     }
 
     public dispose(): void {
@@ -81,8 +113,28 @@ class ManagerReadyImpl implements ManagerReady {
             }
         }
 
-        ids.forEach((managerId) => this.checkExtension(managerId));
-        await Promise.all(Array.from(ids).map((managerId) => this._waitForEnvManager(managerId)));
+        await this.waitForEnvManagerId(Array.from(ids));
+    }
+
+    public async waitForEnvManagerId(managerIds: string[]): Promise<void> {
+        managerIds.forEach((managerId) => this.checkExtension(managerId));
+        await Promise.all(managerIds.map((managerId) => this._waitForEnvManager(managerId)));
+    }
+
+    public async waitForAllEnvManagers(): Promise<void> {
+        const ids: Set<string> = new Set();
+        this.pm.getProjects().forEach((project) => {
+            const m = getDefaultEnvManagerSetting(this.pm, project.uri);
+            if (m && !ids.has(m)) {
+                ids.add(m);
+            }
+        });
+
+        const m = getDefaultEnvManagerSetting(this.pm, undefined);
+        if (m) {
+            ids.add(m);
+        }
+        await this.waitForEnvManagerId(Array.from(ids));
     }
 
     private _waitForPkgManager(managerId: string): Promise<void> {
@@ -111,12 +163,11 @@ class ManagerReadyImpl implements ManagerReady {
             }
         }
 
-        ids.forEach((managerId) => this.checkExtension(managerId));
-        await Promise.all(
-            Array.from(ids).map((managerId) => {
-                return this._waitForPkgManager(managerId);
-            }),
-        );
+        await this.waitForPkgManagerId(Array.from(ids));
+    }
+    public async waitForPkgManagerId(managerIds: string[]): Promise<void> {
+        managerIds.forEach((managerId) => this.checkExtension(managerId));
+        await Promise.all(managerIds.map((managerId) => this._waitForPkgManager(managerId)));
     }
 }
 
@@ -134,7 +185,22 @@ export async function waitForEnvManager(uris?: Uri[]): Promise<void> {
     return mr.waitForEnvManager(uris);
 }
 
+export async function waitForEnvManagerId(managerIds: string[]): Promise<void> {
+    const mr = await _deferred.promise;
+    return mr.waitForEnvManagerId(managerIds);
+}
+
+export async function waitForAllEnvManagers(): Promise<void> {
+    const mr = await _deferred.promise;
+    return mr.waitForAllEnvManagers();
+}
+
 export async function waitForPkgManager(uris?: Uri[]): Promise<void> {
     const mr = await _deferred.promise;
     return mr.waitForPkgManager(uris);
 }
+
+export async function waitForPkgManagerId(managerIds: string[]): Promise<void> {
+    const mr = await _deferred.promise;
+    return mr.waitForPkgManagerId(managerIds);
+}

src/features/envManagers.ts

@@ -10,7 +10,7 @@ import {
     PythonProject,
     SetEnvironmentScope,
 } from '../api';
-import { traceError } from '../common/logging';
+import { traceError, traceVerbose } from '../common/logging';
 import {
     EditAllManagerSettings,
     getDefaultEnvManagerSetting,
@@ -39,7 +39,11 @@ import { sendTelemetryEvent } from '../common/telemetry/sender';
 import { EventNames } from '../common/telemetry/constants';
 
 function generateId(name: string): string {
-    return `${getCallingExtension()}:${name}`;
+    const newName = name.toLowerCase().replace(/[^a-zA-Z0-9-_]/g, '_');
+    if (name !== newName) {
+        traceVerbose(`Environment manager name "${name}"  was normalized to "${newName}"`);
+    }
+    return `${getCallingExtension()}:${newName}`;
 }
 
 export class PythonEnvironmentManagers implements EnvironmentManagers {

src/features/pythonApi.ts

@@ -48,7 +48,7 @@ import { runInTerminal } from './terminal/runInTerminal';
 import { runInBackground } from './execution/runInBackground';
 import { EnvVarManager } from './execution/envVariableManager';
 import { checkUri } from '../common/utils/pathUtils';
-import { waitForEnvManager } from './common/managerReady';
+import { waitForAllEnvManagers, waitForEnvManager, waitForEnvManagerId } from './common/managerReady';
 
 class PythonEnvironmentApiImpl implements PythonEnvironmentApi {
     private readonly _onDidChangeEnvironments = new EventEmitter<DidChangeEnvironmentsEventArgs>();
@@ -163,7 +163,8 @@ class PythonEnvironmentApiImpl implements PythonEnvironmentApi {
             return result;
         }
     }
-    removeEnvironment(environment: PythonEnvironment): Promise<void> {
+    async removeEnvironment(environment: PythonEnvironment): Promise<void> {
+        await waitForEnvManagerId([environment.envId.managerId]);
         const manager = this.envManagers.getEnvironmentManager(environment);
         if (!manager) {
             return Promise.reject(new Error('No environment manager found'));
@@ -174,9 +175,12 @@ class PythonEnvironmentApiImpl implements PythonEnvironmentApi {
         const currentScope = checkUri(scope) as RefreshEnvironmentsScope;
 
         if (currentScope === undefined) {
+            await waitForAllEnvManagers();
             await Promise.all(this.envManagers.managers.map((manager) => manager.refresh(currentScope)));
             return Promise.resolve();
         }
+
+        await waitForEnvManager([currentScope]);
         const manager = this.envManagers.getEnvironmentManager(currentScope);
         if (!manager) {
             return Promise.reject(new Error(`No environment manager found for: ${currentScope.fsPath}`));
@@ -186,10 +190,13 @@ class PythonEnvironmentApiImpl implements PythonEnvironmentApi {
     async getEnvironments(scope: GetEnvironmentsScope): Promise<PythonEnvironment[]> {
         const currentScope = checkUri(scope) as GetEnvironmentsScope;
         if (currentScope === 'all' || currentScope === 'global') {
+            await waitForAllEnvManagers();
             const promises = this.envManagers.managers.map((manager) => manager.getEnvironments(currentScope));
             const items = await Promise.all(promises);
             return items.flat();
         }
+
+        await waitForEnvManager([currentScope]);
         const manager = this.envManagers.getEnvironmentManager(currentScope);
         if (!manager) {
             return [];
@@ -199,14 +206,21 @@ class PythonEnvironmentApiImpl implements PythonEnvironmentApi {
         return items;
     }
     onDidChangeEnvironments: Event<DidChangeEnvironmentsEventArgs> = this._onDidChangeEnvironments.event;
-    setEnvironment(scope: SetEnvironmentScope, environment?: PythonEnvironment): Promise<void> {
-        return this.envManagers.setEnvironment(checkUri(scope) as SetEnvironmentScope, environment);
+    async setEnvironment(scope: SetEnvironmentScope, environment?: PythonEnvironment): Promise<void> {
+        const currentScope = checkUri(scope) as SetEnvironmentScope;
+        await waitForEnvManager(
+            currentScope ? (currentScope instanceof Uri ? [currentScope] : currentScope) : undefined,
+        );
+        return this.envManagers.setEnvironment(currentScope, environment);
     }
     async getEnvironment(scope: GetEnvironmentScope): Promise<PythonEnvironment | undefined> {
-        return this.envManagers.getEnvironment(checkUri(scope) as GetEnvironmentScope);
+        const currentScope = checkUri(scope) as GetEnvironmentScope;
+        await waitForEnvManager(currentScope ? [currentScope] : undefined);
+        return this.envManagers.getEnvironment(currentScope);
     }
     onDidChangeEnvironment: Event<DidChangeEnvironmentEventArgs> = this._onDidChangeEnvironment.event;
     async resolveEnvironment(context: ResolveEnvironmentContext): Promise<PythonEnvironment | undefined> {
+        await waitForAllEnvManagers();
         const projects = this.projectManager.getProjects();
         const projectEnvManagers: InternalEnvironmentManager[] = [];
         projects.forEach((p) => {
@@ -227,21 +241,24 @@ class PythonEnvironmentApiImpl implements PythonEnvironmentApi {
         }
         return new Disposable(() => disposables.forEach((d) => d.dispose()));
     }
-    managePackages(context: PythonEnvironment, options: PackageManagementOptions): Promise<void> {
+    async managePackages(context: PythonEnvironment, options: PackageManagementOptions): Promise<void> {
+        await waitForEnvManagerId([context.envId.managerId]);
         const manager = this.envManagers.getPackageManager(context);
         if (!manager) {
             return Promise.reject(new Error('No package manager found'));
         }
         return manager.manage(context, options);
     }
-    refreshPackages(context: PythonEnvironment): Promise<void> {
+    async refreshPackages(context: PythonEnvironment): Promise<void> {
+        await waitForEnvManagerId([context.envId.managerId]);
         const manager = this.envManagers.getPackageManager(context);
         if (!manager) {
             return Promise.reject(new Error('No package manager found'));
         }
         return manager.refresh(context);
     }
-    getPackages(context: PythonEnvironment): Promise<Package[] | undefined> {
+    async getPackages(context: PythonEnvironment): Promise<Package[] | undefined> {
+        await waitForEnvManagerId([context.envId.managerId]);
         const manager = this.envManagers.getPackageManager(context);
         if (!manager) {
             return Promise.resolve(undefined);