Preselect projects using selected environment

src/common/pickers/projects.ts

@@ -3,11 +3,21 @@ import { QuickPickItem } from 'vscode';
 import { PythonProject } from '../../api';
 import { showQuickPick, showQuickPickWithButtons } from '../window.apis';
 import { Pickers } from '../localize';
+import { normalizePath } from '../utils/pathUtils';
 
 interface ProjectQuickPickItem extends QuickPickItem {
     project: PythonProject;
 }
 
+export interface PickProjectManyOptions {
+    /** Whether to show the Quick Pick back button. */
+    showBackButton?: boolean;
+    /** Whether to show the picker instead of auto-selecting when there is one project. */
+    showPickerForSingleProject?: boolean;
+    /** Projects that should be checked when the picker opens. */
+    selectedProjects?: readonly PythonProject[];
+}
+
 export async function pickProject(projects: ReadonlyArray<PythonProject>): Promise<PythonProject | undefined> {
     if (projects.length > 1) {
         const items: ProjectQuickPickItem[] = projects.map((pw) => ({
@@ -28,11 +38,23 @@ export async function pickProject(projects: ReadonlyArray<PythonProject>): Promi
     return undefined;
 }
 
+/**
+ * Prompts for zero or more projects, optionally preselecting existing project choices.
+ *
+ * The boolean overload is retained for existing callers that only configure back navigation.
+ */
 export async function pickProjectMany(
     projects: readonly PythonProject[],
-    showBackButton?: boolean,
+    showBackButtonOrOptions?: boolean | PickProjectManyOptions,
 ): Promise<PythonProject[] | undefined> {
-    if (projects.length > 1) {
+    const options =
+        typeof showBackButtonOrOptions === 'boolean'
+            ? { showBackButton: showBackButtonOrOptions }
+            : showBackButtonOrOptions;
+    if (projects.length > 1 || (projects.length === 1 && options?.showPickerForSingleProject)) {
+        const selectedProjectPaths = new Set(
+            options?.selectedProjects?.map((project) => normalizePath(path.resolve(project.uri.fsPath))) ?? [],
+        );
         const items: ProjectQuickPickItem[] = projects.map((pw) => ({
             label: path.basename(pw.uri.fsPath),
             description: pw.uri.fsPath,
@@ -42,7 +64,10 @@ export async function pickProjectMany(
             placeHolder: Pickers.Project.selectProjects,
             ignoreFocusOut: true,
             canPickMany: true,
-            showBackButton: showBackButton,
+            showBackButton: options?.showBackButton,
+            selected: items.filter((item) =>
+                selectedProjectPaths.has(normalizePath(path.resolve(item.project.uri.fsPath))),
+            ),
         });
         if (Array.isArray(item)) {
             return item.map((p) => p.project);

src/common/utils/environmentUtils.ts

@@ -0,0 +1,34 @@
+import path from 'path';
+import { PythonEnvironment } from '../../api';
+import { INLINE_SCRIPT_MANAGER_ID } from '../constants';
+import { normalizePath } from './pathUtils';
+
+/**
+ * Determines whether two descriptors should represent the same checked project
+ * in an environment-selection picker.
+ *
+ * Stable IDs take precedence. A normalized absolute interpreter path is used as
+ * a fallback because manager discovery can recreate an equivalent descriptor.
+ */
+export function isSameEnvironmentForProjectSelection(
+    first: PythonEnvironment | undefined,
+    second: PythonEnvironment | undefined,
+): boolean {
+    if (first === second) {
+        return true;
+    }
+    if (!first || !second || first.envId.managerId !== second.envId.managerId) {
+        return false;
+    }
+    if (first.envId.managerId === INLINE_SCRIPT_MANAGER_ID) {
+        return haveSameEnvironmentPath(first, second) && first.version === second.version;
+    }
+    return first.envId.id === second.envId.id || haveSameEnvironmentPath(first, second);
+}
+
+function haveSameEnvironmentPath(first: PythonEnvironment, second: PythonEnvironment): boolean {
+    return (
+        normalizePath(path.resolve(first.environmentPath.fsPath)) ===
+        normalizePath(path.resolve(second.environmentPath.fsPath))
+    );
+}

src/features/envCommands.ts

@@ -49,6 +49,8 @@ import {
     pickWorkspaceFolder,
 } from '../common/pickers/managers';
 import { pickProject, pickProjectMany } from '../common/pickers/projects';
+import { isSameEnvironmentForProjectSelection as isSameEnvironment } from '../common/utils/environmentUtils';
+import { normalizePath } from '../common/utils/pathUtils';
 import { isWindows } from '../common/utils/platformUtils';
 import { handlePythonPath } from '../common/utils/pythonPath';
 import {
@@ -484,10 +486,24 @@ async function setEnvironmentCommandInternal(
             const view = context as PythonEnvTreeItem;
             const projects = wm.getProjects();
             if (projects.length > 0) {
-                const selected = await pickProjectMany(projects);
-                if (selected && selected.length > 0) {
-                    // Check if the selected environment is already the current one for each project
-                    await setEnvironmentForProjects(selected, context.environment, em);
+                const currentEnvironments = await Promise.all(
+                    projects.map((project) => em.getEnvironment(project.uri)),
+                );
+                const projectsUsingEnvironment = projects.filter((_, index) =>
+                    isSameEnvironment(currentEnvironments[index], view.environment),
+                );
+                const selected = await pickProjectMany(projects, {
+                    showPickerForSingleProject: true,
+                    selectedProjects: projectsUsingEnvironment,
+                });
+                if (selected !== undefined) {
+                    await applyProjectEnvironmentSelection(
+                        projects,
+                        projectsUsingEnvironment,
+                        selected,
+                        view.environment,
+                        em,
+                    );
                 }
             } else {
                 await em.setEnvironments('global', view.environment);
@@ -564,10 +580,10 @@ async function setEnvironmentForProjects(
     environment: PythonEnvironment,
     em: EnvironmentManagers,
 ) {
-    let alreadySet: PythonProject[] = [];
+    const alreadySet: PythonProject[] = [];
     for (const p of selectedProjects) {
         const currentEnv = await em.getEnvironment(p.uri);
-        if (currentEnv?.envId.id === environment.envId.id) {
+        if (isSameEnvironment(currentEnv, environment)) {
             alreadySet.push(p);
         }
     }
@@ -587,6 +603,49 @@ async function setEnvironmentForProjects(
     await em.setEnvironments(uris, environment);
 }
 
+async function applyProjectEnvironmentSelection(
+    projects: readonly PythonProject[],
+    previouslySelected: readonly PythonProject[],
+    selected: readonly PythonProject[],
+    environment: PythonEnvironment,
+    em: EnvironmentManagers,
+): Promise<void> {
+    const selectedPaths = new Set(selected.map(getProjectPathKey));
+    const currentEnvironments = await Promise.all(projects.map((project) => em.getEnvironment(project.uri)));
+    const currentEnvironmentByProject = new Map(
+        projects.map((project, index) => [getProjectPathKey(project), currentEnvironments[index]]),
+    );
+    const resetCandidates = previouslySelected.filter(
+        (project) => !selectedPaths.has(getProjectPathKey(project)),
+    );
+    const projectsToSet = selected.filter(
+        (project) => !isSameEnvironment(currentEnvironmentByProject.get(getProjectPathKey(project)), environment),
+    );
+
+    if (projectsToSet.length > 0) {
+        await em.setEnvironments(
+            projectsToSet.map((project) => project.uri),
+            environment,
+        );
+    }
+    const resetEnvironments = await Promise.all(
+        resetCandidates.map((project) => em.getEnvironment(project.uri)),
+    );
+    const projectsToReset = resetCandidates.filter((_, index) =>
+        isSameEnvironment(resetEnvironments[index], environment),
+    );
+    if (projectsToReset.length > 0) {
+        await em.setEnvironments(
+            projectsToReset.map((project) => project.uri),
+            undefined,
+        );
+    }
+}
+
+function getProjectPathKey(project: PythonProject): string {
+    return normalizePath(path.resolve(project.uri.fsPath));
+}
+
 export async function setEnvManagerCommand(em: EnvironmentManagers, wm: PythonProjectManager): Promise<void> {
     const projects = await pickProjectMany(wm.getProjects());
     if (projects && projects.length > 0) {

src/test/common/environmentUtils.unit.test.ts

@@ -0,0 +1,22 @@
+import * as assert from 'assert';
+import * as path from 'path';
+import { isSameEnvironmentForProjectSelection } from '../../common/utils/environmentUtils';
+import { createMockPythonEnvironment } from '../mocks/pythonEnvironment';
+
+suite('Environment identity utilities', () => {
+    test('reconciles same-manager picker descriptors by normalized absolute path', () => {
+        const environmentPath = path.join(process.cwd(), 'same-environment', 'python');
+        const first = createMockPythonEnvironment({
+            envPath: environmentPath,
+            id: 'first-id',
+            managerId: 'ms-python.python:venv',
+        });
+        const second = createMockPythonEnvironment({
+            envPath: environmentPath,
+            id: 'second-id',
+            managerId: 'ms-python.python:venv',
+        });
+
+        assert.strictEqual(isSameEnvironmentForProjectSelection(first, second), true);
+    });
+});

src/test/common/projectPicker.unit.test.ts

@@ -0,0 +1,63 @@
+import * as assert from 'assert';
+import * as path from 'path';
+import * as sinon from 'sinon';
+import { Uri } from 'vscode';
+import { PythonProject } from '../../api';
+import { pickProjectMany } from '../../common/pickers/projects';
+import * as windowApis from '../../common/window.apis';
+
+suite('Project multi-picker', () => {
+    teardown(() => sinon.restore());
+
+    test('preselects the requested projects', async () => {
+        const projects: PythonProject[] = ['one', 'two', 'three'].map((name) => ({
+            name,
+            uri: Uri.file(path.join(process.cwd(), name)),
+        }));
+        const showQuickPick = sinon.stub(windowApis, 'showQuickPickWithButtons').callsFake(async (items, options) => {
+            assert.deepStrictEqual(options?.selected, [items[0], items[2]]);
+            return options?.selected ?? [];
+        });
+
+        const selected = await pickProjectMany(projects, {
+            showBackButton: true,
+            selectedProjects: [projects[0], { ...projects[2] }],
+        });
+
+        assert.deepStrictEqual(selected, [projects[0], projects[2]]);
+        assert.strictEqual(showQuickPick.firstCall.args[1]?.showBackButton, true);
+    });
+
+    test('retains the boolean back-button option for existing callers', async () => {
+        const projects: PythonProject[] = ['one', 'two'].map((name) => ({
+            name,
+            uri: Uri.file(path.join(process.cwd(), name)),
+        }));
+        const showQuickPick = sinon.stub(windowApis, 'showQuickPickWithButtons').resolves(undefined);
+
+        await pickProjectMany(projects, true);
+
+        assert.strictEqual(showQuickPick.firstCall.args[1]?.showBackButton, true);
+        assert.deepStrictEqual(showQuickPick.firstCall.args[1]?.selected, []);
+    });
+
+    test('can show and preselect a single project without changing the legacy shortcut', async () => {
+        const project: PythonProject = {
+            name: 'one',
+            uri: Uri.file(path.join(process.cwd(), 'one')),
+        };
+        const showQuickPick = sinon.stub(windowApis, 'showQuickPickWithButtons').resolves([]);
+
+        const selected = await pickProjectMany([project], {
+            showPickerForSingleProject: true,
+            selectedProjects: [project],
+        });
+
+        assert.deepStrictEqual(selected, []);
+        assert.deepStrictEqual(showQuickPick.firstCall.args[1]?.selected, [showQuickPick.firstCall.args[0][0]]);
+
+        showQuickPick.resetHistory();
+        assert.deepStrictEqual(await pickProjectMany([project]), [project]);
+        sinon.assert.notCalled(showQuickPick);
+    });
+});