Add CSV export option to zoomed in plots (#4252)

Author: julieg18

extension/package.json

@@ -1669,6 +1669,7 @@
     "execa": "5.1.1",
     "fs-extra": "11.1.1",
     "js-yaml": "4.1.0",
+    "json-2-csv": "4.0.0",
     "json5": "2.2.3",
     "lodash.clonedeep": "4.5.0",
     "lodash.get": "4.4.2",

extension/src/fileSystem/index.test.ts

@@ -15,7 +15,8 @@ import {
   isSameOrChild,
   getModifiedTime,
   findOrCreateDvcYamlFile,
-  writeJson
+  writeJson,
+  writeCsv
 } from '.'
 import { dvcDemoPath } from '../test/util'
 import { DOT_DVC } from '../cli/dvc/constants'
@@ -71,6 +72,21 @@ describe('writeJson', () => {
   })
 })
 
+describe('writeCsv', () => {
+  it('should write csv into given file', async () => {
+    await writeCsv('file-name.csv', [
+      { nested: { string: 'string1' }, value: 3 },
+      { nested: { string: 'string2' }, value: 4 },
+      { nested: { string: 'string3' }, value: 6 }
+    ])
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      'file-name.csv',
+      'nested.string,value\nstring1,3\nstring2,4\nstring3,6'
+    )
+  })
+})
+
 describe('findDvcRootPaths', () => {
   it('should find the dvc root if it exists in the given folder', async () => {
     const dvcRoots = await findDvcRootPaths(dvcDemoPath)

extension/src/fileSystem/index.ts

@@ -20,6 +20,7 @@ import {
 } from 'fs-extra'
 import { load } from 'js-yaml'
 import { Uri, workspace, window, commands, ViewColumn } from 'vscode'
+import { json2csv } from 'json-2-csv'
 import { standardizePath } from './path'
 import { definedAndNonEmpty, sortCollectedArray } from '../util/array'
 import { Logger } from '../common/logger'
@@ -208,7 +209,9 @@ export const loadJson = <T>(path: string): T | undefined => {
   }
 }
 
-export const writeJson = <T extends Record<string, unknown>>(
+export const writeJson = <
+  T extends Record<string, unknown> | Array<Record<string, unknown>>
+>(
   path: string,
   obj: T,
   format = false
@@ -218,6 +221,15 @@ export const writeJson = <T extends Record<string, unknown>>(
   return writeFileSync(path, json)
 }
 
+export const writeCsv = async (
+  path: string,
+  arr: Array<Record<string, unknown>>
+) => {
+  ensureFileSync(path)
+  const csv = await json2csv(arr)
+  return writeFileSync(path, csv)
+}
+
 export const getPidFromFile = async (
   path: string
 ): Promise<number | undefined> => {
@@ -269,7 +281,8 @@ export const getBinDisplayText = (
     : path
 }
 
-export const showSaveDialog = (
-  defaultUri: Uri,
-  filters?: { [name: string]: string[] }
-) => window.showSaveDialog({ defaultUri, filters })
+export const showSaveDialog = (fileName: string, extname: string) =>
+  window.showSaveDialog({
+    defaultUri: Uri.file(fileName),
+    filters: { [extname.toUpperCase()]: [extname] }
+  })

extension/src/plots/model/collect.ts

@@ -111,7 +111,7 @@ export const collectCustomPlots = ({
 export const collectCustomPlotRawData = (
   orderValue: CustomPlotsOrderValue,
   experiments: Experiment[]
-) => {
+): Array<Record<string, unknown>> => {
   const { metric, param } = orderValue
   const metricPath = getFullValuePath(ColumnType.METRICS, metric)
   const paramPath = getFullValuePath(ColumnType.PARAMS, param)
@@ -472,7 +472,7 @@ export const collectSelectedTemplatePlotRawData = ({
     multiSourceEncodingUpdate
   )
 
-  return datapoints
+  return datapoints as unknown as Array<Record<string, unknown>>
 }
 
 export const collectOrderedRevisions = (

extension/src/plots/model/index.ts

@@ -57,7 +57,7 @@ import {
 } from '../multiSource/collect'
 import { isDvcError } from '../../cli/dvc/reader'
 import { ErrorsModel } from '../errors/model'
-import { openFileInEditor, writeJson } from '../../fileSystem'
+import { openFileInEditor, writeCsv, writeJson } from '../../fileSystem'
 import { Toast } from '../../vscode/toast'
 
 export class PlotsModel extends ModelWithPersistence {
@@ -227,21 +227,15 @@ export class PlotsModel extends ModelWithPersistence {
     return selectedRevisions
   }
 
-  public savePlotData(plotId: string, filePath: string) {
-    const foundCustomPlot = this.customPlotsOrder.find(
-      ({ metric, param }) => getCustomPlotId(metric, param) === plotId
-    )
-
-    const rawData = foundCustomPlot
-      ? this.getCustomPlotData(foundCustomPlot)
-      : this.getSelectedTemplatePlotData(plotId)
+  public savePlotDataAsJson(filePath: string, plotId: string) {
+    void this.savePlotData(filePath, plotId, data => {
+      writeJson(filePath, data, true)
+      return Promise.resolve()
+    })
+  }
 
-    try {
-      writeJson(filePath, rawData as unknown as Record<string, unknown>, true)
-      void openFileInEditor(filePath)
-    } catch {
-      void Toast.showError('Cannot write to file')
-    }
+  public savePlotDataAsCsv(filePath: string, plotId: string) {
+    void this.savePlotData(filePath, plotId, data => writeCsv(filePath, data))
   }
 
   public getTemplatePlots(
@@ -479,4 +473,25 @@ export class PlotsModel extends ModelWithPersistence {
 
     return collectCustomPlotRawData(orderValue, experiments)
   }
+
+  private async savePlotData(
+    filePath: string,
+    plotId: string,
+    writeToFile: (rawData: Array<Record<string, unknown>>) => Promise<void>
+  ) {
+    const foundCustomPlot = this.customPlotsOrder.find(
+      ({ metric, param }) => getCustomPlotId(metric, param) === plotId
+    )
+
+    const rawData = foundCustomPlot
+      ? this.getCustomPlotData(foundCustomPlot)
+      : this.getSelectedTemplatePlotData(plotId)
+
+    try {
+      await writeToFile(rawData)
+      void openFileInEditor(filePath)
+    } catch {
+      void Toast.showError('Cannot write to file')
+    }
+  }
 }

extension/src/plots/webview/messages.ts

@@ -1,4 +1,4 @@
-import { Uri, commands } from 'vscode'
+import { commands } from 'vscode'
 import isEmpty from 'lodash.isempty'
 import {
   ComparisonPlot,
@@ -82,8 +82,10 @@ export class WebviewMessages {
           RegisteredCommands.PLOTS_CUSTOM_ADD,
           this.dvcRoot
         )
-      case MessageFromWebviewType.EXPORT_PLOT_DATA:
-        return this.exportPlotData(message.payload)
+      case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_CSV:
+        return this.exportPlotDataAsCsv(message.payload)
+      case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_JSON:
+        return this.exportPlotDataAsJson(message.payload)
       case MessageFromWebviewType.RESIZE_PLOTS:
         return this.setPlotSize(
           message.payload.section,
@@ -344,19 +346,35 @@ export class WebviewMessages {
     return this.plots.getCustomPlots() || null
   }
 
-  private async exportPlotData(plotId: string) {
-    const file = await showSaveDialog(Uri.file('data.json'), { JSON: ['json'] })
+  private async exportPlotDataAsJson(plotId: string) {
+    const file = await showSaveDialog('data.json', 'json')
 
     if (!file) {
       return
     }
 
     sendTelemetryEvent(
-      EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA,
+      EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON,
       undefined,
       undefined
     )
 
-    this.plots.savePlotData(plotId, file.path)
+    void this.plots.savePlotDataAsJson(file.path, plotId)
+  }
+
+  private async exportPlotDataAsCsv(plotId: string) {
+    const file = await showSaveDialog('data.csv', 'csv')
+
+    if (!file) {
+      return
+    }
+
+    sendTelemetryEvent(
+      EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV,
+      undefined,
+      undefined
+    )
+
+    void this.plots.savePlotDataAsCsv(file.path, plotId)
   }
 }

extension/src/telemetry/constants.ts

@@ -71,7 +71,8 @@ export const EventName = Object.assign(
       'views.plots.comparisonRowsReordered',
     VIEWS_PLOTS_CREATED: 'views.plots.created',
     VIEWS_PLOTS_EXPERIMENT_TOGGLE: 'views.plots.toggleExperimentStatus',
-    VIEWS_PLOTS_EXPORT_PLOT_DATA: 'views.plots.exportPlotData',
+    VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV: 'views.plots.exportPlotDataAsCsv',
+    VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON: 'views.plots.exportPlotDataAsJson',
     VIEWS_PLOTS_FOCUS_CHANGED: 'views.plots.focusChanged',
     VIEWS_PLOTS_REVISIONS_REORDERED: 'views.plots.revisionsReordered',
     VIEWS_PLOTS_SECTION_RESIZED: 'views.plots.sectionResized',
@@ -267,7 +268,8 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_PLOTS_SELECT_EXPERIMENTS]: undefined
   [EventName.VIEWS_PLOTS_SELECT_PLOTS]: undefined
   [EventName.VIEWS_PLOTS_EXPERIMENT_TOGGLE]: undefined
-  [EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA]: undefined
+  [EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV]: undefined
+  [EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON]: undefined
 
   [EventName.VIEWS_PLOTS_ZOOM_PLOT]: { isImage: boolean }
   [EventName.VIEWS_REORDER_PLOTS_CUSTOM]: undefined