diff --git a/src/common/localize.ts b/src/common/localize.ts index dd7c637f..ab8dc436 100644 --- a/src/common/localize.ts +++ b/src/common/localize.ts @@ -65,6 +65,10 @@ export namespace Pickers { export namespace Project { export const selectProject = l10n.t('Select a project, folder or script'); export const selectProjects = l10n.t('Select one or more projects, folders or scripts'); + export const setForCurrentFile = l10n.t('Set for current file'); + export const addCurrentFileAsProject = l10n.t('Add current file as project...'); + export const currentFileSection = l10n.t('Current File'); + export const projectsSection = l10n.t('Projects'); } export namespace pyProject { diff --git a/src/common/pickers/projects.ts b/src/common/pickers/projects.ts index 19e93f64..c7c53543 100644 --- a/src/common/pickers/projects.ts +++ b/src/common/pickers/projects.ts @@ -1,5 +1,5 @@ import path from 'path'; -import { QuickPickItem } from 'vscode'; +import { QuickPickItem, QuickPickItemKind, Uri } from 'vscode'; import { PythonProject } from '../../api'; import { showQuickPick, showQuickPickWithButtons } from '../window.apis'; import { Pickers } from '../localize'; @@ -8,6 +8,33 @@ interface ProjectQuickPickItem extends QuickPickItem { project: PythonProject; } +export const CURRENT_FILE_ACTION = 'currentFile'; +export const ADD_PROJECT_ACTION = 'addProject'; + +export interface CurrentFileResult { + action: typeof CURRENT_FILE_ACTION; + fileUri: Uri; +} + +export interface AddProjectResult { + action: typeof ADD_PROJECT_ACTION; + fileUri: Uri; +} + +export interface ProjectsResult { + action: 'projects'; + projects: PythonProject[]; +} + +export type ProjectPickerResult = CurrentFileResult | AddProjectResult | ProjectsResult | undefined; + +interface ActionQuickPickItem extends QuickPickItem { + action: typeof CURRENT_FILE_ACTION | typeof ADD_PROJECT_ACTION; + fileUri: Uri; +} + +type EnrichedQuickPickItem = ProjectQuickPickItem | ActionQuickPickItem | QuickPickItem; + export async function pickProject(projects: ReadonlyArray): Promise { if (projects.length > 1) { const items: ProjectQuickPickItem[] = projects.map((pw) => ({ @@ -54,3 +81,72 @@ export async function pickProjectMany( } return undefined; } + +/** + * Shows a project picker with additional "Current File" options at the top. + * When the active editor has a Python file, two special items are injected: + * - "Set for current file" — scopes environment to just the active file URI + * - "Add current file as project..." — creates a project at the file's parent directory + * + * @param projects - The list of existing projects to show + * @param activeFileUri - The URI of the active Python file (if any) + * @returns A discriminated result indicating the user's choice, or undefined if cancelled + */ +export async function pickProjectWithCurrentFile( + projects: readonly PythonProject[], + activeFileUri: Uri, +): Promise { + const items: EnrichedQuickPickItem[] = []; + + // Current file section + items.push({ + label: Pickers.Project.currentFileSection, + kind: QuickPickItemKind.Separator, + }); + items.push({ + label: `$(file) ${Pickers.Project.setForCurrentFile}`, + description: path.basename(activeFileUri.fsPath), + action: CURRENT_FILE_ACTION, + fileUri: activeFileUri, + } as ActionQuickPickItem); + items.push({ + label: `$(add) ${Pickers.Project.addCurrentFileAsProject}`, + description: path.dirname(activeFileUri.fsPath), + action: ADD_PROJECT_ACTION, + fileUri: activeFileUri, + } as ActionQuickPickItem); + + // Projects section + items.push({ + label: Pickers.Project.projectsSection, + kind: QuickPickItemKind.Separator, + }); + for (const pw of projects) { + items.push({ + label: path.basename(pw.uri.fsPath), + description: pw.uri.fsPath, + project: pw, + } as ProjectQuickPickItem); + } + + const selected = await showQuickPickWithButtons(items, { + placeHolder: Pickers.Project.selectProjects, + ignoreFocusOut: true, + }); + + if (!selected) { + return undefined; + } + + if ('action' in selected) { + const actionItem = selected as ActionQuickPickItem; + return { action: actionItem.action, fileUri: actionItem.fileUri }; + } + + if ('project' in selected) { + const projectItem = selected as ProjectQuickPickItem; + return { action: 'projects', projects: [projectItem.project] }; + } + + return undefined; +} diff --git a/src/features/envCommands.ts b/src/features/envCommands.ts index 5c9978d9..f6fc4094 100644 --- a/src/features/envCommands.ts +++ b/src/features/envCommands.ts @@ -40,7 +40,7 @@ import { pickPackageManager, pickWorkspaceFolder, } from '../common/pickers/managers'; -import { pickProject, pickProjectMany } from '../common/pickers/projects'; +import { pickProject, pickProjectMany, pickProjectWithCurrentFile, ADD_PROJECT_ACTION, CURRENT_FILE_ACTION } from '../common/pickers/projects'; import { isWindows } from '../common/utils/platformUtils'; import { handlePythonPath } from '../common/utils/pythonPath'; import { @@ -350,10 +350,41 @@ export async function setEnvironmentCommand( try { const projects = wm.getProjects(); if (projects.length > 0) { - const selected = await pickProjectMany(projects); - if (selected && selected.length > 0) { - const uris = selected.map((p) => p.uri); - await setEnvironmentCommand(uris, em, wm); + // Check if the active editor has a Python file open + const activeEditor = activeTextEditor(); + const activeFileUri = + activeEditor?.document?.languageId === 'python' && + activeEditor.document.uri.scheme === 'file' && + !activeEditor.document.isUntitled + ? activeEditor.document.uri + : undefined; + + if (activeFileUri) { + // Show enriched picker with current file options + const result = await pickProjectWithCurrentFile(projects, activeFileUri); + if (result) { + if (result.action === CURRENT_FILE_ACTION) { + await setEnvironmentCommand([result.fileUri], em, wm); + } else if (result.action === ADD_PROJECT_ACTION) { + const parentUri = Uri.file(path.dirname(result.fileUri.fsPath)); + await executeCommand('python-envs.addPythonProjectGivenResource', parentUri); + // Find the newly created project and open environment picker + const newProject = wm.get(parentUri); + if (newProject) { + await setEnvironmentCommand([newProject.uri], em, wm); + } + } else { + const uris = result.projects.map((p) => p.uri); + await setEnvironmentCommand(uris, em, wm); + } + } + } else { + // No active Python file; use standard multi-select project picker + const selected = await pickProjectMany(projects); + if (selected && selected.length > 0) { + const uris = selected.map((p) => p.uri); + await setEnvironmentCommand(uris, em, wm); + } } } else { const globalEnvManager = em.getEnvironmentManager(undefined); diff --git a/src/test/common/pickers/projects.unit.test.ts b/src/test/common/pickers/projects.unit.test.ts new file mode 100644 index 00000000..b0082509 --- /dev/null +++ b/src/test/common/pickers/projects.unit.test.ts @@ -0,0 +1,139 @@ +import * as assert from 'assert'; +import * as sinon from 'sinon'; +import { QuickPickItemKind, Uri } from 'vscode'; +import { PythonProject } from '../../../api'; +import { Pickers } from '../../../common/localize'; +import { + ADD_PROJECT_ACTION, + CURRENT_FILE_ACTION, + pickProjectWithCurrentFile, +} from '../../../common/pickers/projects'; +import * as windowApis from '../../../common/window.apis'; + +suite('pickProjectWithCurrentFile', () => { + let showQuickPickWithButtonsStub: sinon.SinonStub; + + const project1: PythonProject = { + uri: Uri.file('/workspace/project1'), + name: 'project1', + }; + const project2: PythonProject = { + uri: Uri.file('/workspace/project2'), + name: 'project2', + }; + const activeFileUri = Uri.file('/workspace/project1/main.py'); + + setup(() => { + showQuickPickWithButtonsStub = sinon.stub(windowApis, 'showQuickPickWithButtons'); + }); + + teardown(() => { + sinon.restore(); + }); + + test('should show current file items and project items', async () => { + showQuickPickWithButtonsStub.resolves(undefined); + + await pickProjectWithCurrentFile([project1, project2], activeFileUri); + + assert.ok(showQuickPickWithButtonsStub.calledOnce, 'showQuickPickWithButtons should be called once'); + const items = showQuickPickWithButtonsStub.firstCall.args[0]; + + // Should have: separator + 2 action items + separator + 2 project items = 6 items + assert.strictEqual(items.length, 6, 'Should have 6 items total'); + + // First item: current file separator + assert.strictEqual(items[0].kind, QuickPickItemKind.Separator); + assert.strictEqual(items[0].label, Pickers.Project.currentFileSection); + + // Second item: "Set for current file" + assert.ok(items[1].label.includes(Pickers.Project.setForCurrentFile)); + assert.strictEqual(items[1].action, CURRENT_FILE_ACTION); + assert.strictEqual(items[1].fileUri, activeFileUri); + + // Third item: "Add current file as project..." + assert.ok(items[2].label.includes(Pickers.Project.addCurrentFileAsProject)); + assert.strictEqual(items[2].action, ADD_PROJECT_ACTION); + assert.strictEqual(items[2].fileUri, activeFileUri); + + // Fourth item: projects separator + assert.strictEqual(items[3].kind, QuickPickItemKind.Separator); + assert.strictEqual(items[3].label, Pickers.Project.projectsSection); + + // Fifth and sixth items: projects + assert.strictEqual(items[4].project, project1); + assert.strictEqual(items[5].project, project2); + }); + + test('should return currentFile result when "Set for current file" is selected', async () => { + showQuickPickWithButtonsStub.callsFake((items: unknown[]) => { + // Simulate selecting the "Set for current file" item + return Promise.resolve(items[1]); + }); + + const result = await pickProjectWithCurrentFile([project1], activeFileUri); + + assert.ok(result, 'Result should not be undefined'); + assert.strictEqual(result!.action, CURRENT_FILE_ACTION); + if (result!.action === CURRENT_FILE_ACTION) { + assert.strictEqual(result!.fileUri, activeFileUri); + } + }); + + test('should return addProject result when "Add current file as project..." is selected', async () => { + showQuickPickWithButtonsStub.callsFake((items: unknown[]) => { + // Simulate selecting the "Add current file as project..." item + return Promise.resolve(items[2]); + }); + + const result = await pickProjectWithCurrentFile([project1], activeFileUri); + + assert.ok(result, 'Result should not be undefined'); + assert.strictEqual(result!.action, ADD_PROJECT_ACTION); + if (result!.action === ADD_PROJECT_ACTION) { + assert.strictEqual(result!.fileUri, activeFileUri); + } + }); + + test('should return projects result when a project is selected', async () => { + showQuickPickWithButtonsStub.callsFake((items: unknown[]) => { + // Simulate selecting the first project item (index 4, after separators + action items) + return Promise.resolve(items[4]); + }); + + const result = await pickProjectWithCurrentFile([project1, project2], activeFileUri); + + assert.ok(result, 'Result should not be undefined'); + assert.strictEqual(result!.action, 'projects'); + if (result!.action === 'projects') { + assert.strictEqual(result!.projects.length, 1); + assert.strictEqual(result!.projects[0], project1); + } + }); + + test('should return undefined when picker is cancelled', async () => { + showQuickPickWithButtonsStub.resolves(undefined); + + const result = await pickProjectWithCurrentFile([project1], activeFileUri); + + assert.strictEqual(result, undefined, 'Should return undefined when cancelled'); + }); + + test('should use ignoreFocusOut in picker options', async () => { + showQuickPickWithButtonsStub.resolves(undefined); + + await pickProjectWithCurrentFile([project1], activeFileUri); + + const options = showQuickPickWithButtonsStub.firstCall.args[1]; + assert.strictEqual(options.ignoreFocusOut, true, 'ignoreFocusOut should be true'); + }); + + test('should not use canPickMany', async () => { + showQuickPickWithButtonsStub.resolves(undefined); + + await pickProjectWithCurrentFile([project1], activeFileUri); + + const options = showQuickPickWithButtonsStub.firstCall.args[1]; + assert.strictEqual(options.canPickMany, undefined, 'canPickMany should not be set'); + }); +}); diff --git a/src/test/features/envCommands.unit.test.ts b/src/test/features/envCommands.unit.test.ts index e972d3b4..8b71c2c0 100644 --- a/src/test/features/envCommands.unit.test.ts +++ b/src/test/features/envCommands.unit.test.ts @@ -4,9 +4,11 @@ import * as typeMoq from 'typemoq'; import { Uri } from 'vscode'; import { PythonEnvironment, PythonProject } from '../../api'; import * as commandApi from '../../common/command.api'; +import * as envPickerApi from '../../common/pickers/environments'; import * as managerApi from '../../common/pickers/managers'; import * as projectApi from '../../common/pickers/projects'; -import { createAnyEnvironmentCommand, revealEnvInManagerView } from '../../features/envCommands'; +import * as windowApis from '../../common/window.apis'; +import { createAnyEnvironmentCommand, revealEnvInManagerView, setEnvironmentCommand } from '../../features/envCommands'; import { EnvManagerView } from '../../features/views/envManagersView'; import { ProjectEnvironment, ProjectItem } from '../../features/views/treeViewItems'; import { EnvironmentManagers, InternalEnvironmentManager, PythonProjectManager } from '../../internal.api'; @@ -224,3 +226,225 @@ suite('Reveal Env In Manager View Command Tests', () => { managerView.verify((m) => m.reveal(environment), typeMoq.Times.once()); }); }); + +suite('Set Environment Command - Current File Tests', () => { + let em: typeMoq.IMock; + let wm: typeMoq.IMock; + let manager: typeMoq.IMock; + let env: typeMoq.IMock; + let activeTextEditorStub: sinon.SinonStub; + let pickProjectWithCurrentFileStub: sinon.SinonStub; + let pickProjectManyStub: sinon.SinonStub; + let pickEnvironmentStub: sinon.SinonStub; + let executeCommandStub: sinon.SinonStub; + + const project: PythonProject = { + uri: Uri.file('/workspace/project1'), + name: 'project1', + }; + const activeFileUri = Uri.file('/workspace/project1/main.py'); + + setup(() => { + manager = typeMoq.Mock.ofType(); + manager.setup((m) => m.id).returns(() => 'test'); + manager.setup((m) => m.displayName).returns(() => 'Test Manager'); + + env = typeMoq.Mock.ofType(); + env.setup((e) => e.envId).returns(() => ({ id: 'env1', managerId: 'test' })); + setupNonThenable(env); + + em = typeMoq.Mock.ofType(); + em.setup((e) => e.managers).returns(() => [manager.object]); + em.setup((e) => e.getProjectEnvManagers(typeMoq.It.isAny())).returns(() => [manager.object]); + + wm = typeMoq.Mock.ofType(); + + activeTextEditorStub = sinon.stub(windowApis, 'activeTextEditor'); + pickProjectWithCurrentFileStub = sinon.stub(projectApi, 'pickProjectWithCurrentFile'); + pickProjectManyStub = sinon.stub(projectApi, 'pickProjectMany'); + pickEnvironmentStub = sinon.stub(envPickerApi, 'pickEnvironment'); + executeCommandStub = sinon.stub(commandApi, 'executeCommand'); + }); + + teardown(() => { + sinon.restore(); + }); + + test('should use pickProjectWithCurrentFile when active editor has a Python file', async () => { + // Mock active editor with a Python file + activeTextEditorStub.returns({ + document: { + languageId: 'python', + uri: activeFileUri, + isUntitled: false, + }, + }); + wm.setup((w) => w.getProjects(typeMoq.It.isAny())).returns(() => [project]); + pickProjectWithCurrentFileStub.resolves(undefined); // User cancels + + await setEnvironmentCommand(undefined, em.object, wm.object); + + assert.ok( + pickProjectWithCurrentFileStub.calledOnce, + 'pickProjectWithCurrentFile should be called when Python file is active', + ); + assert.ok(pickProjectManyStub.notCalled, 'pickProjectMany should not be called when Python file is active'); + }); + + test('should use pickProjectMany when no active editor', async () => { + activeTextEditorStub.returns(undefined); + wm.setup((w) => w.getProjects(typeMoq.It.isAny())).returns(() => [project]); + pickProjectManyStub.resolves(undefined); + + await setEnvironmentCommand(undefined, em.object, wm.object); + + assert.ok(pickProjectManyStub.calledOnce, 'pickProjectMany should be called when no active editor'); + assert.ok( + pickProjectWithCurrentFileStub.notCalled, + 'pickProjectWithCurrentFile should not be called when no active editor', + ); + }); + + test('should use pickProjectMany when active editor has non-Python file', async () => { + activeTextEditorStub.returns({ + document: { + languageId: 'javascript', + uri: Uri.file('/workspace/project1/index.js'), + isUntitled: false, + }, + }); + wm.setup((w) => w.getProjects(typeMoq.It.isAny())).returns(() => [project]); + pickProjectManyStub.resolves(undefined); + + await setEnvironmentCommand(undefined, em.object, wm.object); + + assert.ok( + pickProjectManyStub.calledOnce, + 'pickProjectMany should be called for non-Python files', + ); + assert.ok( + pickProjectWithCurrentFileStub.notCalled, + 'pickProjectWithCurrentFile should not be called for non-Python files', + ); + }); + + test('should handle "Set for current file" action by passing file URI', async () => { + activeTextEditorStub.returns({ + document: { + languageId: 'python', + uri: activeFileUri, + isUntitled: false, + }, + }); + wm.setup((w) => w.getProjects(typeMoq.It.isAny())).returns(() => [project]); + + pickProjectWithCurrentFileStub.resolves({ + action: 'currentFile', + fileUri: activeFileUri, + }); + + // The setEnvironmentCommand will be called recursively with [activeFileUri] + // which triggers the Uri[] branch that calls pickEnvironment + manager.setup((m) => m.get(typeMoq.It.isAny())).returns(() => Promise.resolve(undefined)); + pickEnvironmentStub.resolves(env.object); + em.setup((e) => e.setEnvironments(typeMoq.It.isAny(), typeMoq.It.isAny())).returns(() => Promise.resolve()); + + await setEnvironmentCommand(undefined, em.object, wm.object); + + assert.ok(pickEnvironmentStub.calledOnce, 'pickEnvironment should be called after selecting current file'); + }); + + test('should handle "Add current file as project" action', async () => { + activeTextEditorStub.returns({ + document: { + languageId: 'python', + uri: activeFileUri, + isUntitled: false, + }, + }); + wm.setup((w) => w.getProjects(typeMoq.It.isAny())).returns(() => [project]); + + pickProjectWithCurrentFileStub.resolves({ + action: 'addProject', + fileUri: activeFileUri, + }); + + // Mock executeCommand for addPythonProjectGivenResource + executeCommandStub.resolves(); + + // Mock finding the new project after creation + const newProject: PythonProject = { + uri: Uri.file('/workspace/project1'), + name: 'project1', + }; + wm.setup((w) => w.get(typeMoq.It.isAny())).returns(() => newProject); + + // After project is created, the env picker will be shown + manager.setup((m) => m.get(typeMoq.It.isAny())).returns(() => Promise.resolve(undefined)); + pickEnvironmentStub.resolves(env.object); + em.setup((e) => e.setEnvironments(typeMoq.It.isAny(), typeMoq.It.isAny())).returns(() => Promise.resolve()); + + await setEnvironmentCommand(undefined, em.object, wm.object); + + assert.ok( + executeCommandStub.calledWith('python-envs.addPythonProjectGivenResource', sinon.match.any), + 'Should call addPythonProjectGivenResource command', + ); + assert.ok( + pickEnvironmentStub.calledOnce, + 'pickEnvironment should be called after creating the project', + ); + }); + + test('should handle project selection from enriched picker', async () => { + activeTextEditorStub.returns({ + document: { + languageId: 'python', + uri: activeFileUri, + isUntitled: false, + }, + }); + wm.setup((w) => w.getProjects(typeMoq.It.isAny())).returns(() => [project]); + + pickProjectWithCurrentFileStub.resolves({ + action: 'projects', + projects: [project], + }); + + // After project is selected, the env picker will be shown + manager.setup((m) => m.get(typeMoq.It.isAny())).returns(() => Promise.resolve(undefined)); + pickEnvironmentStub.resolves(env.object); + em.setup((e) => e.setEnvironments(typeMoq.It.isAny(), typeMoq.It.isAny())).returns(() => Promise.resolve()); + + await setEnvironmentCommand(undefined, em.object, wm.object); + + assert.ok( + pickEnvironmentStub.calledOnce, + 'pickEnvironment should be called after selecting a project', + ); + }); + + test('should not show current file options when file scheme is not file', async () => { + activeTextEditorStub.returns({ + document: { + languageId: 'python', + uri: Uri.parse('untitled:Untitled-1'), + scheme: 'untitled', + isUntitled: true, + }, + }); + wm.setup((w) => w.getProjects(typeMoq.It.isAny())).returns(() => [project]); + pickProjectManyStub.resolves(undefined); + + await setEnvironmentCommand(undefined, em.object, wm.object); + + assert.ok( + pickProjectManyStub.calledOnce, + 'pickProjectMany should be called for untitled files', + ); + assert.ok( + pickProjectWithCurrentFileStub.notCalled, + 'pickProjectWithCurrentFile should not be called for untitled files', + ); + }); +});