Add "Save as TSV" option to plot modal (#4285)

Author: julieg18

extension/src/fileSystem/index.test.ts

@@ -17,6 +17,7 @@ import {
   findOrCreateDvcYamlFile,
   writeJson,
   writeCsv,
+  writeTsv,
   isPathInProject
 } from '.'
 import { dvcDemoPath } from '../test/util'
@@ -88,6 +89,21 @@ describe('writeCsv', () => {
   })
 })
 
+describe('writeTsv', () => {
+  it('should write tsv into given file', async () => {
+    await writeTsv('file-name.tsv', [
+      { nested: { string: 'string1' }, value: 3 },
+      { nested: { string: 'string2' }, value: 4 },
+      { nested: { string: 'string3' }, value: 6 }
+    ])
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      'file-name.tsv',
+      'nested.string\tvalue\nstring1\t3\nstring2\t4\nstring3\t6'
+    )
+  })
+})
+
 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

@@ -243,6 +243,15 @@ export const writeCsv = async (
   return writeFileSync(path, csv)
 }
 
+export const writeTsv = async (
+  path: string,
+  arr: Array<Record<string, unknown>>
+) => {
+  ensureFileSync(path)
+  const csv = await json2csv(arr, { delimiter: { field: '\t' } })
+  return writeFileSync(path, csv)
+}
+
 export const getPidFromFile = async (
   path: string
 ): Promise<number | undefined> => {

extension/src/plots/model/index.ts

@@ -58,7 +58,12 @@ import {
 } from '../multiSource/collect'
 import { isDvcError } from '../../cli/dvc/reader'
 import { ErrorsModel } from '../errors/model'
-import { openFileInEditor, writeCsv, writeJson } from '../../fileSystem'
+import {
+  openFileInEditor,
+  writeCsv,
+  writeJson,
+  writeTsv
+} from '../../fileSystem'
 import { Toast } from '../../vscode/toast'
 
 export class PlotsModel extends ModelWithPersistence {
@@ -244,6 +249,10 @@ export class PlotsModel extends ModelWithPersistence {
     void this.savePlotData(filePath, plotId, data => writeCsv(filePath, data))
   }
 
+  public savePlotDataAsTsv(filePath: string, plotId: string) {
+    void this.savePlotData(filePath, plotId, data => writeTsv(filePath, data))
+  }
+
   public getTemplatePlots(
     order: TemplateOrder | undefined,
     selectedRevisions: Revision[]

extension/src/plots/webview/messages.ts

@@ -84,6 +84,8 @@ export class WebviewMessages {
         )
       case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_CSV:
         return this.exportPlotDataAsCsv(message.payload)
+      case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_TSV:
+        return this.exportPlotDataAsTsv(message.payload)
       case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_JSON:
         return this.exportPlotDataAsJson(message.payload)
       case MessageFromWebviewType.RESIZE_PLOTS:
@@ -362,35 +364,53 @@ export class WebviewMessages {
     return this.plots.getCustomPlots() || null
   }
 
-  private async exportPlotDataAsJson(plotId: string) {
-    const file = await showSaveDialog('data.json', 'json')
+  private async exportPlotData(
+    extName: string,
+    plotId: string,
+    event:
+      | typeof EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV
+      | typeof EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON
+      | typeof EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_TSV,
+    writeFile: (filePath: string, plotId: string) => void
+  ) {
+    const file = await showSaveDialog(`data.${extName}`, extName)
 
     if (!file) {
       return
     }
 
-    sendTelemetryEvent(
-      EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON,
-      undefined,
-      undefined
-    )
+    sendTelemetryEvent(event, undefined, undefined)
 
-    void this.plots.savePlotDataAsJson(file.path, plotId)
+    writeFile(file.path, plotId)
   }
 
-  private async exportPlotDataAsCsv(plotId: string) {
-    const file = await showSaveDialog('data.csv', 'csv')
-
-    if (!file) {
-      return
-    }
+  private exportPlotDataAsJson(plotId: string) {
+    void this.exportPlotData(
+      'json',
+      plotId,
+      EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON,
+      (filePath: string, plotId: string) =>
+        this.plots.savePlotDataAsJson(filePath, plotId)
+    )
+  }
 
-    sendTelemetryEvent(
+  private exportPlotDataAsCsv(plotId: string) {
+    void this.exportPlotData(
+      'csv',
+      plotId,
       EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV,
-      undefined,
-      undefined
+      (filePath: string, plotId: string) =>
+        this.plots.savePlotDataAsCsv(filePath, plotId)
     )
+  }
 
-    void this.plots.savePlotDataAsCsv(file.path, plotId)
+  private exportPlotDataAsTsv(plotId: string) {
+    void this.exportPlotData(
+      'tsv',
+      plotId,
+      EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_TSV,
+      (filePath: string, plotId: string) =>
+        this.plots.savePlotDataAsTsv(filePath, plotId)
+    )
   }
 }

extension/src/telemetry/constants.ts

@@ -73,6 +73,7 @@ export const EventName = Object.assign(
     VIEWS_PLOTS_EXPERIMENT_TOGGLE: 'views.plots.toggleExperimentStatus',
     VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV: 'views.plots.exportPlotDataAsCsv',
     VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON: 'views.plots.exportPlotDataAsJson',
+    VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_TSV: 'views.plots.exportPlotDataAsTsv',
     VIEWS_PLOTS_FOCUS_CHANGED: 'views.plots.focusChanged',
     VIEWS_PLOTS_REVISIONS_REORDERED: 'views.plots.revisionsReordered',
     VIEWS_PLOTS_SECTION_RESIZED: 'views.plots.sectionResized',
@@ -271,6 +272,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_PLOTS_EXPERIMENT_TOGGLE]: undefined
   [EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV]: undefined
   [EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON]: undefined
+  [EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_TSV]: undefined
 
   [EventName.VIEWS_PLOTS_ZOOM_PLOT]: { isImage: boolean }
   [EventName.VIEWS_REORDER_PLOTS_CUSTOM]: undefined

extension/src/test/suite/plots/index.test.ts

@@ -489,14 +489,13 @@ suite('Plots Test Suite', () => {
       const mockWriteJson = stub(FileSystem, 'writeJson')
       const mockOpenFile = stub(FileSystem, 'openFileInEditor')
       const exportFile = Uri.file('raw-data.json')
-      const mockShowInformationMessage = stub(window, 'showErrorMessage')
 
       mockShowSaveDialog.resolves(exportFile)
 
-      const fileCancelledEvent = new Promise(resolve =>
-        mockShowSaveDialog.onFirstCall().callsFake(() => {
+      const openFileEvent = new Promise(resolve =>
+        mockOpenFile.onFirstCall().callsFake(() => {
           resolve(undefined)
-          return Promise.resolve(undefined)
+          return Promise.resolve(undefined as unknown as TextDocument)
         })
       )
 
@@ -505,27 +504,84 @@ suite('Plots Test Suite', () => {
         type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_JSON
       })
 
-      await fileCancelledEvent
+      await openFileEvent
 
-      expect(mockSendTelemetryEvent).not.to.be.called
-      expect(mockWriteJson).not.to.be.called
+      expect(mockWriteJson).to.be.calledOnce
+      expect(mockWriteJson).to.be.calledWithExactly(
+        exportFile.path,
+        customPlot.values,
+        true
+      )
+      expect(mockOpenFile).to.calledWithExactly(exportFile.path)
+      expect(mockSendTelemetryEvent).to.be.calledOnce
+      expect(mockSendTelemetryEvent).to.be.calledWithExactly(
+        EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON,
+        undefined,
+        undefined
+      )
+    })
 
-      const jsonWriteErrorEvent = new Promise(resolve =>
-        mockWriteJson.onFirstCall().callsFake(() => {
+    it('should handle an export plot data as csv message from the webview', async () => {
+      const { plots } = await buildPlots({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+
+      const webview = await plots.showWebview()
+      const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+      const mockShowSaveDialog = stub(window, 'showSaveDialog')
+      const mockWriteCsv = stub(FileSystem, 'writeCsv')
+      const mockOpenFile = stub(FileSystem, 'openFileInEditor')
+      const exportFile = Uri.file('raw-data.csv')
+      const templatePlot = templatePlotsFixture.plots[0].entries[0]
+
+      mockShowSaveDialog.resolves(exportFile)
+
+      const openFileEvent = new Promise(resolve =>
+        mockOpenFile.onFirstCall().callsFake(() => {
           resolve(undefined)
-          throw new Error('file failed to write')
+          return Promise.resolve(undefined as unknown as TextDocument)
         })
       )
 
       mockMessageReceived.fire({
-        payload: customPlot.id,
-        type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_JSON
+        payload: templatePlot.id,
+        type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_CSV
       })
 
-      await jsonWriteErrorEvent
+      await openFileEvent
 
-      expect(mockOpenFile).not.to.be.called
-      expect(mockShowInformationMessage).to.be.called
+      expect(mockWriteCsv).to.be.calledOnce
+      expect(mockWriteCsv).to.be.calledWithExactly(
+        exportFile.path,
+        (templatePlot.content.data as { values: unknown[] }).values
+      )
+      expect(mockOpenFile).to.calledWithExactly(exportFile.path)
+      expect(mockSendTelemetryEvent).to.be.calledOnce
+      expect(mockSendTelemetryEvent).to.be.calledWithExactly(
+        EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV,
+        undefined,
+        undefined
+      )
+    })
+
+    it('should handle an export plot data as tsv message from the webview', async () => {
+      const { plots } = await buildPlots({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+
+      const webview = await plots.showWebview()
+      const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+      const mockShowSaveDialog = stub(window, 'showSaveDialog')
+      const mockWriteTsv = stub(FileSystem, 'writeTsv')
+      const mockOpenFile = stub(FileSystem, 'openFileInEditor')
+      const exportFile = Uri.file('raw-data.tsv')
+      const customPlot = customPlotsFixture.plots[0]
+
+      mockShowSaveDialog.resolves(exportFile)
 
       const openFileEvent = new Promise(resolve =>
         mockOpenFile.onFirstCall().callsFake(() => {
@@ -536,25 +592,26 @@ suite('Plots Test Suite', () => {
 
       mockMessageReceived.fire({
         payload: customPlot.id,
-        type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_JSON
+        type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_TSV
       })
 
       await openFileEvent
 
-      expect(mockWriteJson).to.be.calledWithExactly(
+      expect(mockWriteTsv).to.be.calledOnce
+      expect(mockWriteTsv).to.be.calledWithExactly(
         exportFile.path,
-        customPlot.values,
-        true
+        customPlot.values
       )
       expect(mockOpenFile).to.calledWithExactly(exportFile.path)
+      expect(mockSendTelemetryEvent).to.be.calledOnce
       expect(mockSendTelemetryEvent).to.be.calledWithExactly(
-        EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_JSON,
+        EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_TSV,
         undefined,
         undefined
       )
     })
 
-    it('should handle an export plot data as csv message from the webview', async () => {
+    it('should handle export data messages from the webview when the file is cancelled or errors are thrown during file writing', async () => {
       const { plots } = await buildPlots({
         disposer: disposable,
         plotsDiff: plotsDiffFixture
@@ -589,10 +646,11 @@ suite('Plots Test Suite', () => {
       expect(mockWriteCsv).not.to.be.called
       expect(mockOpenFile).not.to.be.called
 
-      const openFileEvent = new Promise(resolve =>
-        mockOpenFile.onFirstCall().callsFake(() => {
+      const mockShowInformationMessage = stub(window, 'showErrorMessage')
+      const fileWriteErrorEvent = new Promise(resolve =>
+        mockWriteCsv.onFirstCall().callsFake(() => {
           resolve(undefined)
-          return Promise.resolve(undefined as unknown as TextDocument)
+          throw new Error('file failed to write')
         })
       )
 
@@ -601,18 +659,10 @@ suite('Plots Test Suite', () => {
         type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_CSV
       })
 
-      await openFileEvent
+      await fileWriteErrorEvent
 
-      expect(mockWriteCsv).to.be.calledWithExactly(
-        exportFile.path,
-        (templatePlot.content.data as { values: unknown[] }).values
-      )
-      expect(mockOpenFile).to.calledWithExactly(exportFile.path)
-      expect(mockSendTelemetryEvent).to.be.calledWithExactly(
-        EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA_AS_CSV,
-        undefined,
-        undefined
-      )
+      expect(mockOpenFile).not.to.be.called
+      expect(mockShowInformationMessage).to.be.called
     })
 
     it('should handle a custom plots reordered message from the webview', async () => {

extension/src/webview/contract.ts

@@ -20,6 +20,7 @@ export enum MessageFromWebviewType {
   CREATE_BRANCH_FROM_EXPERIMENT = 'create-branch-from-experiment',
   EXPORT_PLOT_DATA_AS_JSON = 'export-plot-data-as-json',
   EXPORT_PLOT_DATA_AS_CSV = 'export-plot-data-as-csv',
+  EXPORT_PLOT_DATA_AS_TSV = 'export-plot-data-as-tsv',
   FOCUS_FILTERS_TREE = 'focus-filters-tree',
   FOCUS_SORTS_TREE = 'focus-sorts-tree',
   OPEN_EXPERIMENTS_WEBVIEW = 'open-experiments-webview',
@@ -108,6 +109,10 @@ export type MessageFromWebview =
       type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_CSV
       payload: string
     }
+  | {
+      type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_TSV
+      payload: string
+    }
   | {
       type: MessageFromWebviewType.REORDER_COLUMNS
       payload: string[]