Add top-level plots wizard

extension/package.json

@@ -423,6 +423,12 @@
         "category": "DVC",
         "icon": "$(symbol-class)"
       },
+      {
+        "title": "Add Plot",
+        "command": "dvc.addTopLevelPlot",
+        "category": "DVC",
+        "icon": "$(add)"
+      },
       {
         "title": "Show Plots",
         "command": "dvc.showPlots",
@@ -903,6 +909,10 @@
           "command": "dvc.showPipelineDAG",
           "when": "dvc.commands.available && dvc.project.available"
         },
+        {
+          "command": "dvc.addTopLevelPlot",
+          "when": "dvc.commands.available && dvc.project.available"
+        },
         {
           "command": "dvc.showExperiments",
           "when": "dvc.commands.available && dvc.project.available"
@@ -1414,6 +1424,11 @@
           "when": "view == dvc.views.experimentsFilterByTree && dvc.experiments.filtered",
           "group": "navigation@3"
         },
+        {
+          "command": "dvc.addTopLevelPlot",
+          "when": "view == dvc.views.plotsPathsTree",
+          "group": "navigation@0"
+        },
         {
           "command": "dvc.showPlots",
           "when": "view == dvc.views.plotsPathsTree",
@@ -1704,7 +1719,6 @@
     "appdirs": "1.1.0",
     "execa": "5.1.1",
     "fs-extra": "11.1.1",
-    "js-yaml": "4.1.0",
     "json-2-csv": "4.1.0",
     "json5": "2.2.3",
     "lodash.clonedeep": "4.5.0",
@@ -1718,7 +1732,8 @@
     "tree-kill": "1.2.2",
     "uuid": "9.0.0",
     "vega-util": "1.17.2",
-    "vscode-languageclient": "8.1.0"
+    "vscode-languageclient": "8.1.0",
+    "yaml": "2.3.2"
   },
   "devDependencies": {
     "@swc/core": "1.3.82",
@@ -1728,7 +1743,6 @@
     "@types/copy-webpack-plugin": "10.1.0",
     "@types/fs-extra": "11.0.1",
     "@types/jest": "29.5.4",
-    "@types/js-yaml": "4.0.5",
     "@types/lodash.clonedeep": "4.5.7",
     "@types/lodash.get": "4.4.7",
     "@types/lodash.isempty": "4.4.7",

extension/src/cli/dvc/contract.ts

@@ -3,6 +3,18 @@ import { Plot } from '../../plots/webview/contract'
 export const MIN_CLI_VERSION = '2.58.1'
 export const LATEST_TESTED_CLI_VERSION = '3.19.0'
 
+export const PLOT_TEMPLATES = [
+  'simple',
+  'linear',
+  'confusion',
+  'confusion_normalized',
+  'scatter',
+  'scatter_jitter',
+  'smooth',
+  'bar_horizontal_sorted',
+  'bar_horizontal'
+]
+
 type ErrorContents = { type: string; msg: string }
 
 export type DvcError = { error: ErrorContents }

extension/src/commands/external.ts

@@ -72,6 +72,7 @@ export enum RegisteredCommands {
   STOP_EXPERIMENTS = 'dvc.stopAllRunningExperiments',
 
   PIPELINE_SHOW_DAG = 'dvc.showPipelineDAG',
+  PIPELINE_ADD_PLOT = 'dvc.addTopLevelPlot',
 
   PLOTS_PATH_TOGGLE = 'dvc.views.plotsPathsTree.toggleStatus',
   PLOTS_SHOW = 'dvc.showPlots',

extension/src/fileSystem/index.test.ts

@@ -21,7 +21,9 @@ import {
   writeTsv,
   isPathInProject,
   getPidFromFile,
-  getEntryFromJsonFile
+  getEntryFromJsonFile,
+  addPlotToDvcYamlFile,
+  loadDataFile
 } from '.'
 import { dvcDemoPath } from '../test/util'
 import { DOT_DVC } from '../cli/dvc/constants'
@@ -61,6 +63,83 @@ beforeEach(() => {
   jest.resetAllMocks()
 })
 
+describe('loadDataFile', () => {
+  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')
+
+    expect(result).toStrictEqual([
+      { acc: 0.69, epoch: 10 },
+      { acc: 0.345, epoch: 11 }
+    ])
+  })
+
+  it('should load in json file contents', async () => {
+    const mockJsonContent = JSON.stringify([
+      { acc: 0.69, epoch: 10 },
+      { acc: 0.345, epoch: 11 }
+    ])
+
+    mockedReadFileSync.mockReturnValueOnce(mockJsonContent)
+
+    const result = await loadDataFile('values.json')
+
+    expect(result).toStrictEqual([
+      { acc: 0.69, epoch: 10 },
+      { acc: 0.345, epoch: 11 }
+    ])
+  })
+
+  it('should load in tsv file contents', async () => {
+    const mockTsvContent = ['epoch\tacc', '10\t0.69', '11\t0.345'].join('\n')
+
+    mockedReadFileSync.mockReturnValueOnce(mockTsvContent)
+
+    const result = await loadDataFile('values.tsv')
+
+    expect(result).toStrictEqual([
+      { acc: 0.69, epoch: 10 },
+      { acc: 0.345, epoch: 11 }
+    ])
+  })
+
+  it('should load in yaml file contents', async () => {
+    const mockYamlContent = [
+      'stages:',
+      '  train:',
+      '    cmd: python train.py'
+    ].join('\n')
+
+    mockedReadFileSync.mockReturnValueOnce(mockYamlContent)
+
+    const result = await loadDataFile('dvc.yaml')
+
+    expect(result).toStrictEqual({
+      stages: {
+        train: {
+          cmd: 'python train.py'
+        }
+      }
+    })
+  })
+
+  it('should catch any errors thrown during file parsing', async () => {
+    const dataFiles = ['values.csv', 'file.json', 'file.tsv', 'dvc.yaml']
+    mockedReadFileSync.mockImplementation(() => {
+      throw new Error('fake error')
+    })
+
+    for (const file of dataFiles) {
+      const resultWithErr = await loadDataFile(file)
+
+      expect(resultWithErr).toStrictEqual(undefined)
+    }
+  })
+})
+
 describe('writeJson', () => {
   it('should write unformatted json in given file', () => {
     writeJson('file-name.json', { array: [1, 2, 3], number: 1 })
@@ -436,6 +515,124 @@ describe('findOrCreateDvcYamlFile', () => {
   })
 })
 
+describe('addPlotToDvcYamlFile', () => {
+  const mockStagesLines = ['stages:', '  train:', '    cmd: python train.py']
+  const mockPlotsListLines = [
+    'plots:',
+    '  - eval/importance.png',
+    '  - Precision-Recall:',
+    '      x: recall',
+    '      y:',
+    '        eval/prc/train.json: precision',
+    '        eval/prc/test.json: precision'
+  ]
+  const mockNewPlotLines = [
+    '  - data.json:',
+    '      template: simple',
+    '      x: epochs',
+    '      y: accuracy'
+  ]
+  it('should add a plots list with the new plot if the dvc.yaml file has no plots', () => {
+    const mockDvcYamlContent = mockStagesLines.join('\n')
+    const mockPlotYamlContent = ['', 'plots:', ...mockNewPlotLines, ''].join(
+      '\n'
+    )
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+
+    addPlotToDvcYamlFile('/', {
+      dataFile: '/data.json',
+      template: 'simple',
+      x: 'epochs',
+      y: 'accuracy'
+    })
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      '//dvc.yaml',
+      mockDvcYamlContent + mockPlotYamlContent
+    )
+  })
+
+  it('should add the new plot if the dvc.yaml file already has plots', () => {
+    const mockDvcYamlContent = [...mockPlotsListLines, ...mockStagesLines]
+    const mockPlotYamlContent = [...mockNewPlotLines, '']
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent.join('\n'))
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent.join('\n'))
+
+    addPlotToDvcYamlFile('/', {
+      dataFile: '/data.json',
+      template: 'simple',
+      x: 'epochs',
+      y: 'accuracy'
+    })
+
+    mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent)
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      '//dvc.yaml',
+      mockDvcYamlContent.join('\n')
+    )
+  })
+
+  it('should add a new plot if the dvc.yaml plots list is at bottom of file', () => {
+    const mockDvcYamlContent = [...mockStagesLines, ...mockPlotsListLines].join(
+      '\n'
+    )
+    const mockPlotYamlContent = ['', ...mockNewPlotLines, ''].join('\n')
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+
+    addPlotToDvcYamlFile('/', {
+      dataFile: '/data.json',
+      template: 'simple',
+      x: 'epochs',
+      y: 'accuracy'
+    })
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      '//dvc.yaml',
+      mockDvcYamlContent + mockPlotYamlContent
+    )
+  })
+
+  it('should add a new plot with an indent level that matches the dvc.yaml file', () => {
+    const mockDvcYamlContent = [
+      'stages:',
+      '    train:',
+      '        cmd: python train.py',
+      'plots:',
+      '    - eval/importance.png',
+      '    - Precision-Recall:',
+      '          x: recall',
+      '          y:',
+      '              eval/prc/train.json: precision',
+      '              eval/prc/test.json: precision'
+    ].join('\n')
+    const mockPlotYamlContent = [
+      '',
+      '    - data.json:',
+      '          template: simple',
+      '          x: epochs',
+      '          y: accuracy',
+      ''
+    ].join('\n')
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+
+    addPlotToDvcYamlFile('/', {
+      dataFile: '/data.json',
+      template: 'simple',
+      x: 'epochs',
+      y: 'accuracy'
+    })
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      '//dvc.yaml',
+      mockDvcYamlContent + mockPlotYamlContent
+    )
+  })
+})
+
 describe('isPathInProject', () => {
   it('should return true if the path is in the project', () => {
     const path = join(dvcDemoPath, 'dvc.yaml')