Add Multiple X Field Selection to Plot Wizard (#4797)

Author: julieg18

extension/src/fileSystem/index.test.ts

@@ -625,8 +625,8 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'data.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -654,8 +654,38 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'simple_plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'acc.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'acc.json': ['accuracy'] }
+    })
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      '//dvc.yaml',
+      mockDvcYamlContent + mockPlotYamlContent
+    )
+  })
+
+  it('should add the new plot with an axis having multiple fields', () => {
+    const mockDvcYamlContent = mockStagesLines.join('\n')
+    const mockPlotYamlContent = [
+      '',
+      'plots:',
+      '  - simple_plot:',
+      '      template: simple',
+      '      x: epochs',
+      '      y:',
+      '        data.json:',
+      '          - accuracy',
+      '          - epochs',
+      ''
+    ].join('\n')
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+
+    addPlotToDvcYamlFile('/', {
+      template: 'simple',
+      title: 'simple_plot',
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy', 'epochs'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -673,8 +703,8 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'data.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
     })
 
     mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent)
@@ -696,8 +726,8 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'data.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -734,8 +764,8 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'simple_plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'data.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(

extension/src/fileSystem/index.ts

@@ -30,7 +30,7 @@ import { processExists } from '../process/execution'
 import { getFirstWorkspaceFolder } from '../vscode/workspaceFolders'
 import { DOT_DVC } from '../cli/dvc/constants'
 import { delay } from '../util/time'
-import { PlotConfigData } from '../pipeline/quickPick'
+import { PlotConfigData, PlotConfigDataAxis } from '../pipeline/quickPick'
 
 export const exists = (path: string): boolean => existsSync(path)
 
@@ -214,18 +214,35 @@ const loadYamlAsDoc = (
   }
 }
 
+const formatPlotYamlObjAxis = (axis: PlotConfigDataAxis) => {
+  const formattedAxis: { [file: string]: string | string[] } = {}
+
+  for (const [file, fields] of Object.entries(axis)) {
+    if (fields.length === 1) {
+      formattedAxis[file] = fields[0]
+      continue
+    }
+
+    formattedAxis[file] = fields
+  }
+
+  return formattedAxis
+}
+
 const getPlotYamlObj = (plot: PlotConfigData) => {
   const { x, y, template, title } = plot
 
   const yFiles = Object.keys(y)
-  const [xFile, xKey] = Object.entries(x)[0]
-  const oneFileUsed = yFiles.length === 1 && yFiles[0] === xFile
+  const xFiles = Object.keys(x)
+  const firstXFile = xFiles[0]
+  const oneFileUsed =
+    yFiles.length === 1 && xFiles.length === 1 && yFiles[0] === firstXFile
 
   return {
     [title]: {
       template,
-      x: oneFileUsed ? xKey : x,
-      y
+      x: oneFileUsed ? x[firstXFile][0] : formatPlotYamlObjAxis(x),
+      y: formatPlotYamlObjAxis(y)
     }
   }
 }