diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts index 25730afe61..65a5874bc4 100644 --- a/extension/src/fileSystem/index.test.ts +++ b/extension/src/fileSystem/index.test.ts @@ -23,7 +23,7 @@ import { getPidFromFile, getEntryFromJsonFile, addPlotToDvcYamlFile, - loadDataFile + loadDataFiles } from '.' import { dvcDemoPath } from '../test/util' import { DOT_DVC } from '../cli/dvc/constants' @@ -63,17 +63,22 @@ beforeEach(() => { jest.resetAllMocks() }) -describe('loadDataFile', () => { +describe('loadDataFiles', () => { it('should load in csv file contents', async () => { const mockCsvContent = ['epoch,acc', '10,0.69', '11,0.345'].join('\n') mockedReadFileSync.mockReturnValueOnce(mockCsvContent) - const result = await loadDataFile('values.csv') + const result = await loadDataFiles(['values.csv']) expect(result).toStrictEqual([ - { acc: 0.69, epoch: 10 }, - { acc: 0.345, epoch: 11 } + { + data: [ + { acc: 0.69, epoch: 10 }, + { acc: 0.345, epoch: 11 } + ], + file: 'values.csv' + } ]) }) @@ -85,11 +90,16 @@ describe('loadDataFile', () => { mockedReadFileSync.mockReturnValueOnce(mockJsonContent) - const result = await loadDataFile('values.json') + const result = await loadDataFiles(['values.json']) expect(result).toStrictEqual([ - { acc: 0.69, epoch: 10 }, - { acc: 0.345, epoch: 11 } + { + data: [ + { acc: 0.69, epoch: 10 }, + { acc: 0.345, epoch: 11 } + ], + file: 'values.json' + } ]) }) @@ -98,11 +108,16 @@ describe('loadDataFile', () => { mockedReadFileSync.mockReturnValueOnce(mockTsvContent) - const result = await loadDataFile('values.tsv') + const result = await loadDataFiles(['values.tsv']) expect(result).toStrictEqual([ - { acc: 0.69, epoch: 10 }, - { acc: 0.345, epoch: 11 } + { + data: [ + { acc: 0.69, epoch: 10 }, + { acc: 0.345, epoch: 11 } + ], + file: 'values.tsv' + } ]) }) @@ -115,15 +130,47 @@ describe('loadDataFile', () => { mockedReadFileSync.mockReturnValueOnce(mockYamlContent) - const result = await loadDataFile('dvc.yaml') + const result = await loadDataFiles(['dvc.yaml']) - expect(result).toStrictEqual({ - stages: { - train: { - cmd: 'python train.py' - } + expect(result).toStrictEqual([ + { + data: { + stages: { + train: { + cmd: 'python train.py' + } + } + }, + file: 'dvc.yaml' } - }) + ]) + }) + + it('should load in the contents of multiple files', async () => { + const mockTsvContent = ['epoch\tacc', '10\t0.69', '11\t0.345'].join('\n') + const mockCsvContent = ['epoch2,acc2', '10,0.679', '11,0.3'].join('\n') + + mockedReadFileSync.mockReturnValueOnce(mockTsvContent) + mockedReadFileSync.mockReturnValueOnce(mockCsvContent) + + const result = await loadDataFiles(['values.tsv', 'values2.csv']) + + expect(result).toStrictEqual([ + { + data: [ + { acc: 0.69, epoch: 10 }, + { acc: 0.345, epoch: 11 } + ], + file: 'values.tsv' + }, + { + data: [ + { acc2: 0.679, epoch2: 10 }, + { acc2: 0.3, epoch2: 11 } + ], + file: 'values2.csv' + } + ]) }) it('should catch any errors thrown during file parsing', async () => { @@ -133,11 +180,29 @@ describe('loadDataFile', () => { }) for (const file of dataFiles) { - const resultWithErr = await loadDataFile(file) + const resultWithErr = await loadDataFiles([file]) expect(resultWithErr).toStrictEqual(undefined) } }) + + it('should catch any errors thrown during the parsing of multiple files', async () => { + const dataFiles = ['values.csv', 'file.tsv', 'file.json'] + const mockCsvContent = ['epoch,acc', '10,0.69', '11,0.345'].join('\n') + const mockJsonContent = JSON.stringify([ + { acc: 0.69, epoch: 10 }, + { acc: 0.345, epoch: 11 } + ]) + mockedReadFileSync + .mockReturnValueOnce(mockCsvContent) + .mockImplementationOnce(() => { + throw new Error('fake error') + }) + .mockReturnValueOnce(mockJsonContent) + + const resultWithErr = await loadDataFiles(dataFiles) + expect(resultWithErr).toStrictEqual(undefined) + }) }) describe('writeJson', () => { @@ -527,10 +592,11 @@ describe('addPlotToDvcYamlFile', () => { ' eval/prc/test.json: precision' ] const mockNewPlotLines = [ - ' - data.json:', + ' - simple_plot:', ' template: simple', ' x: epochs', - ' y: accuracy' + ' y:', + ' data.json: accuracy' ] it('should add a plots list with the new plot if the dvc.yaml file has no plots', () => { const mockDvcYamlContent = mockStagesLines.join('\n') @@ -541,10 +607,37 @@ describe('addPlotToDvcYamlFile', () => { mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent) addPlotToDvcYamlFile('/', { - dataFile: '/data.json', template: 'simple', - x: 'epochs', - y: 'accuracy' + x: { file: '/data.json', key: 'epochs' }, + y: { file: '/data.json', key: 'accuracy' } + }) + + expect(mockedWriteFileSync).toHaveBeenCalledWith( + '//dvc.yaml', + mockDvcYamlContent + mockPlotYamlContent + ) + }) + + it('should add the new plot with fields coming from different files', () => { + const mockDvcYamlContent = mockStagesLines.join('\n') + const mockPlotYamlContent = [ + '', + 'plots:', + ' - simple_plot:', + ' template: simple', + ' x:', + ' data.json: epochs', + ' y:', + ' acc.json: accuracy', + '' + ].join('\n') + mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent) + mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent) + + addPlotToDvcYamlFile('/', { + template: 'simple', + x: { file: '/data.json', key: 'epochs' }, + y: { file: '/acc.json', key: 'accuracy' } }) expect(mockedWriteFileSync).toHaveBeenCalledWith( @@ -560,10 +653,9 @@ describe('addPlotToDvcYamlFile', () => { mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent.join('\n')) addPlotToDvcYamlFile('/', { - dataFile: '/data.json', template: 'simple', - x: 'epochs', - y: 'accuracy' + x: { file: '/data.json', key: 'epochs' }, + y: { file: '/data.json', key: 'accuracy' } }) mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent) @@ -583,10 +675,9 @@ describe('addPlotToDvcYamlFile', () => { mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent) addPlotToDvcYamlFile('/', { - dataFile: '/data.json', template: 'simple', - x: 'epochs', - y: 'accuracy' + x: { file: '/data.json', key: 'epochs' }, + y: { file: '/data.json', key: 'accuracy' } }) expect(mockedWriteFileSync).toHaveBeenCalledWith( @@ -610,20 +701,20 @@ describe('addPlotToDvcYamlFile', () => { ].join('\n') const mockPlotYamlContent = [ '', - ' - data.json:', + ' - simple_plot:', ' template: simple', ' x: epochs', - ' y: accuracy', + ' y:', + ' data.json: accuracy', '' ].join('\n') mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent) mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent) addPlotToDvcYamlFile('/', { - dataFile: '/data.json', template: 'simple', - x: 'epochs', - y: 'accuracy' + x: { file: '/data.json', key: 'epochs' }, + y: { file: '/data.json', key: 'accuracy' } }) expect(mockedWriteFileSync).toHaveBeenCalledWith( diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts index 2ddfee9289..f79dcb6211 100644 --- a/extension/src/fileSystem/index.ts +++ b/extension/src/fileSystem/index.ts @@ -214,13 +214,23 @@ const loadYamlAsDoc = ( } } +const getPlotYamlObj = (cwd: string, plot: PlotConfigData) => { + const { x, y, template } = plot + const plotName = `${template}_plot` + return { + [plotName]: { + template, + x: x.file === y.file ? x.key : { [relative(cwd, x.file)]: x.key }, + y: { [relative(cwd, y.file)]: y.key } + } + } +} + const getPlotsYaml = ( cwd: string, plotObj: PlotConfigData, indentSearchLines: string[] ) => { - const { dataFile, ...plot } = plotObj - const plotName = relative(cwd, dataFile) const indentReg = /^( +)[^ ]/ const indentLine = indentSearchLines.find(line => indentReg.test(line)) || '' const spacesMatches = indentLine.match(indentReg) @@ -228,7 +238,7 @@ const getPlotsYaml = ( return yaml .stringify( - { plots: [{ [plotName]: plot }] }, + { plots: [getPlotYamlObj(cwd, plotObj)] }, { indent: spaces ? spaces.length : 2 } ) .split('\n') @@ -315,7 +325,7 @@ const loadTsv = (path: string) => { } } -export const loadDataFile = (file: string): unknown => { +const loadDataFile = (file: string): unknown => { const ext = getFileExtension(file) switch (ext) { @@ -330,6 +340,22 @@ export const loadDataFile = (file: string): unknown => { } } +export const loadDataFiles = async ( + files: string[] +): Promise<{ file: string; data: unknown }[] | undefined> => { + const filesData: { file: string; data: unknown }[] = [] + for (const file of files) { + const data = await loadDataFile(file) + + if (!data) { + return undefined + } + + filesData.push({ data, file }) + } + return filesData +} + export const writeJson = < T extends Record | Array> >( diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts index d69f0944f8..ae019ad728 100644 --- a/extension/src/pipeline/index.ts +++ b/extension/src/pipeline/index.ts @@ -119,7 +119,7 @@ export class Pipeline extends DeferredDisposable { return } - const plotConfiguration = await pickPlotConfiguration() + const plotConfiguration = await pickPlotConfiguration(cwd) if (!plotConfiguration) { return diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts index e43baab160..05bf97bfd2 100644 --- a/extension/src/pipeline/quickPick.test.ts +++ b/extension/src/pipeline/quickPick.test.ts @@ -1,14 +1,16 @@ +import { QuickPickItemKind } from 'vscode' import { pickPlotConfiguration } from './quickPick' -import { pickFile } from '../vscode/resourcePicker' -import { quickPickOne } from '../vscode/quickPick' -import { getFileExtension, loadDataFile } from '../fileSystem' +import { pickFiles } from '../vscode/resourcePicker' +import { quickPickOne, quickPickValue } from '../vscode/quickPick' +import { getFileExtension, loadDataFiles } from '../fileSystem' import { Title } from '../vscode/title' import { Toast } from '../vscode/toast' -const mockedPickFile = jest.mocked(pickFile) -const mockedLoadDataFile = jest.mocked(loadDataFile) +const mockedPickFiles = jest.mocked(pickFiles) +const mockedLoadDataFiles = jest.mocked(loadDataFiles) const mockedGetFileExt = jest.mocked(getFileExtension) const mockedQuickPickOne = jest.mocked(quickPickOne) +const mockedQuickPickValue = jest.mocked(quickPickValue) const mockedToast = jest.mocked(Toast) const mockedShowError = jest.fn() mockedToast.showError = mockedShowError @@ -47,39 +49,51 @@ const mockValidData = [ describe('pickPlotConfiguration', () => { it('should let the user pick from files with accepted data types', async () => { - mockedPickFile.mockResolvedValueOnce('file.json') - mockedLoadDataFile.mockReturnValueOnce(mockValidData) + mockedPickFiles.mockResolvedValueOnce(['file.json']) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: mockValidData, file: 'file.json' } + ]) - await pickPlotConfiguration() + await pickPlotConfiguration('/') - expect(mockedPickFile).toHaveBeenCalledWith(Title.SELECT_PLOT_DATA, { + expect(mockedPickFiles).toHaveBeenCalledWith(Title.SELECT_PLOT_DATA, { 'Data Formats': ['json', 'csv', 'tsv', 'yaml'] }) }) it('should return early if user does not select a file', async () => { - mockedPickFile.mockResolvedValueOnce(undefined) + mockedPickFiles.mockResolvedValueOnce(undefined) - const result = await pickPlotConfiguration() + const result = await pickPlotConfiguration('/') expect(result).toStrictEqual(undefined) }) - it('should show a toast message if the file fails to parse', async () => { - mockedPickFile.mockResolvedValueOnce('file.csv') - mockedLoadDataFile.mockReturnValueOnce(undefined) + it('should show a toast message if the files are not the same data type', async () => { + mockedPickFiles.mockResolvedValueOnce(['file.json', 'file.csv']) - const result = await pickPlotConfiguration() + const result = await pickPlotConfiguration('/') + + expect(result).toStrictEqual(undefined) + expect(mockedShowError).toHaveBeenCalledTimes(1) + expect(mockedShowError).toHaveBeenCalledWith('Files must of the same type.') + }) + + it('should show a toast message if the file or files fail to parse', async () => { + mockedPickFiles.mockResolvedValueOnce(['data.csv']) + mockedLoadDataFiles.mockResolvedValueOnce(undefined) + + const result = await pickPlotConfiguration('/') expect(result).toStrictEqual(undefined) expect(mockedShowError).toHaveBeenCalledTimes(1) expect(mockedShowError).toHaveBeenCalledWith( - 'Failed to parse the requested file. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' + 'Failed to parse the requested file or files. Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' ) }) - it('should show a toast message if an array of objects (with atleast two keys) are not found within a file', async () => { - mockedPickFile.mockResolvedValue('file.yaml') + it('should show a toast message if an array of objects (with atleast two keys) are not found within a single chosen file', async () => { + mockedPickFiles.mockResolvedValue(['file.yaml']) const invalidValues: unknown[] = [ 'string', 13, @@ -115,20 +129,43 @@ describe('pickPlotConfiguration', () => { let result for (const [ind, invalidVal] of invalidValues.entries()) { - mockedLoadDataFile.mockReturnValueOnce(invalidVal) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: invalidVal, file: 'file.yaml' } + ]) - result = await pickPlotConfiguration() + result = await pickPlotConfiguration('/') expect(result).toStrictEqual(undefined) expect(mockedShowError).toHaveBeenCalledTimes(1 + ind) expect(mockedShowError).toHaveBeenCalledWith( - 'The requested file does not contain enough keys (columns) to generate a plot. Does the file follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' + 'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' ) } }) - it('should parse fields from valid data files', async () => { - mockedPickFile.mockResolvedValue('file.yaml') + it('should show a toast message if an array of objects (with atleast one key) are not found within multiple chosen files', async () => { + mockedPickFiles.mockResolvedValueOnce([ + 'file.yaml', + 'file2.yaml', + 'file3.yaml' + ]) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' }, + { data: [], file: 'file3.yaml' }, + { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' } + ]) + + const result = await pickPlotConfiguration('/') + + expect(result).toStrictEqual(undefined) + expect(mockedShowError).toHaveBeenCalledTimes(1) + expect(mockedShowError).toHaveBeenCalledWith( + 'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' + ) + }) + + it('should parse fields from a valid data file', async () => { + mockedPickFiles.mockResolvedValue(['file.yaml']) const validValues: unknown[] = [ [{ field1: 1, field2: 2 }], [ @@ -145,24 +182,45 @@ describe('pickPlotConfiguration', () => { ] for (const [ind, val] of validValues.entries()) { - mockedLoadDataFile.mockReturnValueOnce(val) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: val, file: 'file.yaml' } + ]) - await pickPlotConfiguration() + await pickPlotConfiguration('/') expect(mockedShowError).not.toHaveBeenCalled() expect(mockedQuickPickOne).toHaveBeenCalledTimes(ind + 1) } }) + it('should parse fields from multiple valid data files (if atleast two fields are found total)', async () => { + mockedPickFiles.mockResolvedValueOnce([ + 'file.yaml', + 'file2.yaml', + 'file3.yaml' + ]) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' }, + { data: { val: [{ field1: 1 }] }, file: 'file2.yaml' } + ]) + + const result = await pickPlotConfiguration('/') + + expect(result).toStrictEqual(undefined) + expect(mockedShowError).not.toHaveBeenCalled() + }) + it('should let the user pick a template, x field, and y field', async () => { - mockedPickFile.mockResolvedValueOnce('file.json') - mockedLoadDataFile.mockReturnValueOnce(mockValidData) - mockedQuickPickOne - .mockResolvedValueOnce('simple') - .mockResolvedValueOnce('actual') - .mockResolvedValueOnce('prob') + mockedPickFiles.mockResolvedValueOnce(['/file.json']) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: mockValidData, file: '/file.json' } + ]) + mockedQuickPickOne.mockResolvedValueOnce('simple') + mockedQuickPickValue + .mockResolvedValueOnce({ file: '/file.json', key: 'actual' }) + .mockResolvedValueOnce({ file: '/file.json', key: 'prob' }) - const result = await pickPlotConfiguration() + const result = await pickPlotConfiguration('/') expect(mockedQuickPickOne).toHaveBeenNthCalledWith( 1, @@ -179,30 +237,109 @@ describe('pickPlotConfiguration', () => { ], 'Pick a Plot Template' ) - expect(mockedQuickPickOne).toHaveBeenNthCalledWith( + expect(mockedQuickPickValue).toHaveBeenNthCalledWith( + 1, + [ + { + kind: QuickPickItemKind.Separator, + label: 'file.json', + value: undefined + }, + { label: 'actual', value: { file: '/file.json', key: 'actual' } }, + { label: 'prob', value: { file: '/file.json', key: 'prob' } } + ], + { title: Title.SELECT_PLOT_X_METRIC } + ) + expect(mockedQuickPickValue).toHaveBeenNthCalledWith( 2, - ['actual', 'prob'], - 'Pick a Metric for X' + [ + { + kind: QuickPickItemKind.Separator, + label: 'file.json', + value: undefined + }, + { label: 'prob', value: { file: '/file.json', key: 'prob' } } + ], + { + title: Title.SELECT_PLOT_Y_METRIC + } ) - expect(mockedQuickPickOne).toHaveBeenNthCalledWith( - 3, - ['prob'], - 'Pick a Metric for Y' + expect(result).toStrictEqual({ + template: 'simple', + x: { file: '/file.json', key: 'actual' }, + y: { file: '/file.json', key: 'prob' } + }) + }) + + it('should let the user pick a x field and y field from multiple files', async () => { + mockedPickFiles.mockResolvedValueOnce(['/file.json', '/file2.json']) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: mockValidData, file: '/file.json' }, + { data: mockValidData, file: '/file2.json' } + ]) + mockedQuickPickOne.mockResolvedValueOnce('simple') + mockedQuickPickValue + .mockResolvedValueOnce({ file: '/file.json', key: 'actual' }) + .mockResolvedValueOnce({ file: '/file2.json', key: 'prob' }) + + const result = await pickPlotConfiguration('/') + + expect(mockedQuickPickValue).toHaveBeenNthCalledWith( + 1, + [ + { + kind: QuickPickItemKind.Separator, + label: 'file.json', + value: undefined + }, + { label: 'actual', value: { file: '/file.json', key: 'actual' } }, + { label: 'prob', value: { file: '/file.json', key: 'prob' } }, + { + kind: QuickPickItemKind.Separator, + label: 'file2.json', + value: undefined + }, + { label: 'actual', value: { file: '/file2.json', key: 'actual' } }, + { label: 'prob', value: { file: '/file2.json', key: 'prob' } } + ], + { title: Title.SELECT_PLOT_X_METRIC } + ) + expect(mockedQuickPickValue).toHaveBeenNthCalledWith( + 2, + [ + { + kind: QuickPickItemKind.Separator, + label: 'file.json', + value: undefined + }, + { label: 'prob', value: { file: '/file.json', key: 'prob' } }, + { + kind: QuickPickItemKind.Separator, + label: 'file2.json', + value: undefined + }, + { label: 'actual', value: { file: '/file2.json', key: 'actual' } }, + { label: 'prob', value: { file: '/file2.json', key: 'prob' } } + ], + { + title: Title.SELECT_PLOT_Y_METRIC + } ) expect(result).toStrictEqual({ - dataFile: 'file.json', template: 'simple', - x: 'actual', - y: 'prob' + x: { file: '/file.json', key: 'actual' }, + y: { file: '/file2.json', key: 'prob' } }) }) it('should return early if the user does not pick a template', async () => { - mockedPickFile.mockResolvedValueOnce('file.json') - mockedLoadDataFile.mockReturnValueOnce(mockValidData) + mockedPickFiles.mockResolvedValueOnce(['file.json']) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: mockValidData, file: 'file.json' } + ]) mockedQuickPickOne.mockResolvedValueOnce(undefined) - const result = await pickPlotConfiguration() + const result = await pickPlotConfiguration('/') expect(mockedQuickPickOne).toHaveBeenCalledTimes(1) @@ -210,29 +347,32 @@ describe('pickPlotConfiguration', () => { }) it('should return early if the user does not pick a x field', async () => { - mockedPickFile.mockResolvedValueOnce('file.json') - mockedLoadDataFile.mockReturnValueOnce(mockValidData) - mockedQuickPickOne - .mockResolvedValueOnce('simple') - .mockResolvedValueOnce(undefined) + mockedPickFiles.mockResolvedValueOnce(['file.json']) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: mockValidData, file: 'file.json' } + ]) + mockedQuickPickOne.mockResolvedValueOnce('simple') + mockedQuickPickValue.mockResolvedValueOnce(undefined) - const result = await pickPlotConfiguration() + const result = await pickPlotConfiguration('/') - expect(mockedQuickPickOne).toHaveBeenCalledTimes(2) + expect(mockedQuickPickValue).toHaveBeenCalledTimes(1) expect(result).toStrictEqual(undefined) }) it('should return early if the user does not pick a y field', async () => { - mockedPickFile.mockResolvedValueOnce('file.json') - mockedLoadDataFile.mockReturnValueOnce(mockValidData) - mockedQuickPickOne - .mockResolvedValueOnce('simple') + mockedPickFiles.mockResolvedValueOnce(['file.json']) + mockedLoadDataFiles.mockResolvedValueOnce([ + { data: mockValidData, file: 'file.json' } + ]) + mockedQuickPickOne.mockResolvedValueOnce('simple') + mockedQuickPickValue .mockResolvedValueOnce('actual') .mockResolvedValueOnce(undefined) - const result = await pickPlotConfiguration() + const result = await pickPlotConfiguration('/') - expect(mockedQuickPickOne).toHaveBeenCalledTimes(3) + expect(mockedQuickPickValue).toHaveBeenCalledTimes(2) expect(result).toStrictEqual(undefined) }) }) diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts index fd03adf667..0a2d994431 100644 --- a/extension/src/pipeline/quickPick.ts +++ b/extension/src/pipeline/quickPick.ts @@ -1,40 +1,57 @@ +import { relative } from 'path' import isEqual from 'lodash.isequal' +import { QuickPickItemKind } from 'vscode' import { PLOT_TEMPLATES, Value, ValueTree, isValueTree } from '../cli/dvc/contract' -import { loadDataFile } from '../fileSystem' -import { quickPickOne } from '../vscode/quickPick' -import { pickFile } from '../vscode/resourcePicker' +import { getFileExtension, loadDataFiles } from '../fileSystem' +import { quickPickOne, quickPickValue } from '../vscode/quickPick' +import { pickFiles } from '../vscode/resourcePicker' import { Title } from '../vscode/title' import { Toast } from '../vscode/toast' -const pickDataFile = () => { - return pickFile(Title.SELECT_PLOT_DATA, { +const pickDataFiles = (): Thenable => + pickFiles(Title.SELECT_PLOT_DATA, { 'Data Formats': ['json', 'csv', 'tsv', 'yaml'] }) -} const pickTemplateAndFields = async ( - fields: string[] -): Promise<{ x: string; y: string; template: string } | undefined> => { + cwd: string, + fields: { + [file: string]: string[] + } +): Promise => { const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template') if (!template) { return } - const x = await quickPickOne(fields, 'Pick a Metric for X') + const items = [] + + for (const [file, keys] of Object.entries(fields)) { + items.push( + { + kind: QuickPickItemKind.Separator, + label: relative(cwd, file), + value: undefined + }, + ...keys.map(key => ({ label: key, value: { file, key } })) + ) + } + + const x = await quickPickValue(items, { title: Title.SELECT_PLOT_X_METRIC }) if (!x) { return } - const y = await quickPickOne( - fields.filter(field => x !== field), - 'Pick a Metric for Y' + const y = await quickPickValue( + items.filter(item => !isEqual(item.value, x)), + { title: Title.SELECT_PLOT_Y_METRIC } ) if (!y) { @@ -45,10 +62,9 @@ const pickTemplateAndFields = async ( } export type PlotConfigData = { - dataFile: string + x: { file: string; key: string } template: string - x: string - y: string + y: { file: string; key: string } } type UnknownValue = Value | ValueTree @@ -65,7 +81,7 @@ const getFieldsFromArr = (dataArr: UnknownValue[]) => { return objsHaveSameKeys ? fieldObjKeys : [] } -const getFieldOptions = (data: UnknownValue): string[] => { +const getFieldsFromValue = (data: UnknownValue): string[] => { const isArray = Array.isArray(data) const isObj = isValueTree(data) if (!isArray && !isObj) { @@ -79,36 +95,79 @@ const getFieldOptions = (data: UnknownValue): string[] => { : [] } -export const pickPlotConfiguration = async (): Promise< - PlotConfigData | undefined -> => { - const file = await pickDataFile() +const showNotEnoughKeysToast = () => + Toast.showError( + 'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' + ) + +const getFieldsFromDataFiles = ( + dataArr: { data: UnknownValue; file: string }[] +) => { + const keys: { + [file: string]: string[] + } = {} + let keysAmount = 0 + + for (const { file, data } of dataArr) { + const fields = getFieldsFromValue(data) + + if (fields.length === 0) { + void showNotEnoughKeysToast() + return + } + + keysAmount += fields.length + + if (!keys[file]) { + keys[file] = [] + } + keys[file].push(...fields) + } + + if (keysAmount < 2) { + void showNotEnoughKeysToast() + return + } + + return keys +} + +export const pickPlotConfiguration = async ( + cwd: string +): Promise => { + const files = await pickDataFiles() - if (!file) { + if (!files) { return } - const data = await loadDataFile(file) + const fileExts = new Set(files.map(file => getFileExtension(file))) - if (!data) { - return Toast.showError( - 'Failed to parse the requested file. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' - ) + if (fileExts.size > 1) { + return Toast.showError('Files must of the same type.') } - const keys = getFieldOptions(data as UnknownValue) + const filesData = await loadDataFiles(files) - if (keys.length < 2) { + if (!filesData) { return Toast.showError( - 'The requested file does not contain enough keys (columns) to generate a plot. Does the file follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' + 'Failed to parse the requested file or files. Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?' ) } - const templateAndFields = await pickTemplateAndFields(keys) + const keys = getFieldsFromDataFiles( + filesData as { data: UnknownValue; file: string }[] + ) + + if (!keys) { + return + } + + const templateAndFields = await pickTemplateAndFields(cwd, keys) if (!templateAndFields) { return } - return { ...templateAndFields, dataFile: file } + return { ...templateAndFields } } diff --git a/extension/src/plots/quickPick.test.ts b/extension/src/plots/quickPick.test.ts index 7e22b95dee..8bcaa48235 100644 --- a/extension/src/plots/quickPick.test.ts +++ b/extension/src/plots/quickPick.test.ts @@ -20,7 +20,7 @@ describe('pickPlotType', () => { [ { description: - 'Create a data series plot definition by selecting data from a file', + 'Create a data series plot definition by selecting data from one or more files (To select multiple files, hold the Ctrl/Command key and click on the files.).', label: 'Data Series', value: 'data-series' }, diff --git a/extension/src/plots/quickPick.ts b/extension/src/plots/quickPick.ts index 5c00508247..e2358bb262 100644 --- a/extension/src/plots/quickPick.ts +++ b/extension/src/plots/quickPick.ts @@ -11,7 +11,7 @@ export const pickPlotType = () => [ { description: - 'Create a data series plot definition by selecting data from a file', + 'Create a data series plot definition by selecting data from one or more files (To select multiple files, hold the Ctrl/Command key and click on the files.).', label: 'Data Series', value: PLOT_TYPE.DATA_SERIES }, diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts index bde79ed3ff..9ba5a92915 100644 --- a/extension/src/test/suite/pipeline/index.test.ts +++ b/extension/src/test/suite/pipeline/index.test.ts @@ -310,10 +310,9 @@ suite('Pipeline Test Suite', () => { expect(mockAddPlotToDvcFile).not.to.be.called mockPickPlotConfiguration.onSecondCall().resolves({ - dataFile: 'results.json', template: 'simple', - x: 'step', - y: 'acc' + x: { file: 'results.json', key: 'step' }, + y: { file: 'results.json', key: 'acc' } }) await pipeline.addDataSeriesPlot() diff --git a/extension/src/vscode/resourcePicker.test.ts b/extension/src/vscode/resourcePicker.test.ts index 23256ffc5f..69050eb2dc 100644 --- a/extension/src/vscode/resourcePicker.test.ts +++ b/extension/src/vscode/resourcePicker.test.ts @@ -1,6 +1,6 @@ import { resolve } from 'path' import { Uri, window } from 'vscode' -import { pickFile, pickResources } from './resourcePicker' +import { pickFile, pickResources, pickFiles } from './resourcePicker' import { Title } from './title' jest.mock('vscode') @@ -21,6 +21,7 @@ describe('pickFile', () => { await pickFile(mockedTitle) expect(mockedShowOpenDialog).toHaveBeenCalledWith({ + canSelectFiles: true, canSelectFolders: false, canSelectMany: false, openLabel: 'Select', @@ -67,3 +68,34 @@ describe('pickResources', () => { expect(pickedResources).toStrictEqual([mockedUri]) }) }) + +describe('pickFiles', () => { + it('should call window.showOpenDialog with the correct options', async () => { + const mockedTitle = 'I decided to not decide' as Title + mockedShowOpenDialog.mockResolvedValueOnce(undefined) + + await pickFiles(mockedTitle, { Images: ['png', 'jpg'] }) + + expect(mockedShowOpenDialog).toHaveBeenCalledWith({ + canSelectFiles: true, + canSelectFolders: false, + canSelectMany: true, + filters: { Images: ['png', 'jpg'] }, + openLabel: 'Select', + title: mockedTitle + }) + }) + + it('should return an array of paths if any are selected', async () => { + const mockedUris = [ + Uri.file(resolve('mock', 'file1.json')), + Uri.file(resolve('mock', 'file2.json')) + ] + const mockedTitle = 'this is even more fun' as Title + mockedShowOpenDialog.mockResolvedValueOnce(mockedUris) + + const pickedResources = await pickFiles(mockedTitle) + + expect(pickedResources).toStrictEqual(mockedUris.map(uri => uri.fsPath)) + }) +}) diff --git a/extension/src/vscode/resourcePicker.ts b/extension/src/vscode/resourcePicker.ts index 43b6b5aa15..87991ee764 100644 --- a/extension/src/vscode/resourcePicker.ts +++ b/extension/src/vscode/resourcePicker.ts @@ -1,14 +1,16 @@ import { OpenDialogOptions, Uri, window } from 'vscode' import { Title } from './title' -export const pickFile = async ( +export const pickResources = ( title: Title, + canSelectFolders = true, + canSelectMany = true, filters?: OpenDialogOptions['filters'] -): Promise => { +): Thenable => { const opts: OpenDialogOptions = { - canSelectFolders: false, - canSelectMany: false, - filters, + canSelectFiles: true, + canSelectFolders, + canSelectMany, openLabel: 'Select', title } @@ -17,20 +19,24 @@ export const pickFile = async ( opts.filters = filters } - const uris = await window.showOpenDialog(opts) + return window.showOpenDialog(opts) +} +export const pickFile = async (title: Title): Promise => { + const uris = await pickResources(title, false, false) if (uris) { const [{ fsPath }] = uris return fsPath } } -export const pickResources = (title: Title): Thenable => { - return window.showOpenDialog({ - canSelectFiles: true, - canSelectFolders: true, - canSelectMany: true, - openLabel: 'Select', - title - }) +export const pickFiles = async ( + title: Title, + filters?: OpenDialogOptions['filters'] +): Promise => { + const uris = await pickResources(title, false, true, filters) + + if (uris) { + return uris.map(({ fsPath }) => fsPath) + } } diff --git a/extension/src/vscode/title.ts b/extension/src/vscode/title.ts index d9fbbf486b..1545c25e83 100644 --- a/extension/src/vscode/title.ts +++ b/extension/src/vscode/title.ts @@ -40,6 +40,8 @@ export enum Title { SELECT_SORT_DIRECTION = 'Select Sort Direction', SELECT_SORTS_TO_REMOVE = 'Select Sort(s) to Remove', SELECT_TRAINING_SCRIPT = 'Select your Training Script', + SELECT_PLOT_X_METRIC = 'Select a Metric for X', + SELECT_PLOT_Y_METRIC = 'Select a Metric for Y', SELECT_PLOT_DATA = 'Select Data to Plot', SELECT_PLOT_TYPE = 'Select Plot Type', SETUP_WORKSPACE = 'Setup the Workspace',