Add "Export Raw Data" option to plot popup (#4181)

Author: julieg18

extension/src/fileSystem/index.test.ts

@@ -1,5 +1,11 @@
 import { join, relative, resolve } from 'path'
-import { appendFileSync, ensureDirSync, ensureFileSync, remove } from 'fs-extra'
+import {
+  appendFileSync,
+  ensureDirSync,
+  ensureFileSync,
+  remove,
+  writeFileSync
+} from 'fs-extra'
 import { TextDocument, window, workspace } from 'vscode'
 import {
   exists,
@@ -8,7 +14,8 @@ import {
   isDirectory,
   isSameOrChild,
   getModifiedTime,
-  findOrCreateDvcYamlFile
+  findOrCreateDvcYamlFile,
+  writeJson
 } from '.'
 import { dvcDemoPath } from '../test/util'
 import { DOT_DVC } from '../cli/dvc/constants'
@@ -21,12 +28,14 @@ jest.mock('fs-extra', () => {
     __esModule: true,
     ...actualModule,
     appendFileSync: jest.fn(),
-    ensureFileSync: jest.fn()
+    ensureFileSync: jest.fn(),
+    writeFileSync: jest.fn()
   }
 })
 
 const mockedAppendFileSync = jest.mocked(appendFileSync)
 const mockedEnsureFileSync = jest.mocked(ensureFileSync)
+const mockedWriteFileSync = jest.mocked(writeFileSync)
 const mockedWorkspace = jest.mocked(workspace)
 const mockedWindow = jest.mocked(window)
 const mockedOpenTextDocument = jest.fn()
@@ -39,6 +48,29 @@ beforeEach(() => {
   jest.resetAllMocks()
 })
 
+describe('writeJson', () => {
+  it('should write unformatted json in given file', () => {
+    writeJson('file-name.json', { array: [1, 2, 3], number: 1 })
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      'file-name.json',
+      '{"array":[1,2,3],"number":1}'
+    )
+  })
+
+  it('should write formatted json in given file', () => {
+    writeJson('file-name.json', { array: [1, 2, 3], number: 1 }, true)
+
+    const formattedJson =
+      '{\n    "array": [\n        1,\n        2,\n        3\n    ],\n    "number": 1\n}'
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      'file-name.json',
+      formattedJson
+    )
+  })
+})
+
 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

@@ -210,10 +210,12 @@ export const loadJson = <T>(path: string): T | undefined => {
 
 export const writeJson = <T extends Record<string, unknown>>(
   path: string,
-  obj: T
+  obj: T,
+  format = false
 ): void => {
   ensureFileSync(path)
-  return writeFileSync(path, JSON.stringify(obj))
+  const json = format ? JSON.stringify(obj, null, 4) : JSON.stringify(obj)
+  return writeFileSync(path, json)
 }
 
 export const getPidFromFile = async (
@@ -266,3 +268,8 @@ export const getBinDisplayText = (
     ? '.' + sep + relative(workspaceRoot, path)
     : path
 }
+
+export const showSaveDialog = (
+  defaultUri: Uri,
+  filters?: { [name: string]: string[] }
+) => window.showSaveDialog({ defaultUri, filters })

extension/src/plots/model/collect.ts

@@ -108,6 +108,17 @@ export const collectCustomPlots = ({
   return plots
 }
 
+export const collectCustomPlotRawData = (
+  orderValue: CustomPlotsOrderValue,
+  experiments: Experiment[]
+) => {
+  const { metric, param } = orderValue
+  const metricPath = getFullValuePath(ColumnType.METRICS, metric)
+  const paramPath = getFullValuePath(ColumnType.PARAMS, param)
+
+  return getValues(experiments, metricPath, paramPath)
+}
+
 type RevisionPathData = { [path: string]: Record<string, unknown>[] }
 
 export type RevisionData = {
@@ -433,9 +444,37 @@ export const collectSelectedTemplatePlots = (
       group
     })
   }
+
   return acc.length > 0 ? acc : undefined
 }
 
+export const collectSelectedTemplatePlotRawData = ({
+  selectedRevisions,
+  path,
+  template,
+  revisionData,
+  multiSourceEncodingUpdate
+}: {
+  selectedRevisions: string[]
+  path: string
+  template: string
+  revisionData: RevisionData
+  multiSourceEncodingUpdate: { strokeDash: StrokeDashEncoding }
+}) => {
+  const isMultiView = isMultiViewPlot(
+    JSON.parse(template) as TopLevelSpec | VisualizationSpec
+  )
+  const { datapoints } = transformRevisionData(
+    path,
+    selectedRevisions,
+    revisionData,
+    isMultiView,
+    multiSourceEncodingUpdate
+  )
+
+  return datapoints
+}
+
 export const collectOrderedRevisions = (
   revisions: {
     id: string

extension/src/plots/model/index.ts

@@ -10,7 +10,9 @@ import {
   getCustomPlotId,
   collectOrderedRevisions,
   collectImageUrl,
-  collectIdShas
+  collectIdShas,
+  collectSelectedTemplatePlotRawData,
+  collectCustomPlotRawData
 } from './collect'
 import { getRevisionSummaryColumns } from './util'
 import {
@@ -55,6 +57,8 @@ import {
 } from '../multiSource/collect'
 import { isDvcError } from '../../cli/dvc/reader'
 import { ErrorsModel } from '../errors/model'
+import { openFileInEditor, writeJson } from '../../fileSystem'
+import { Toast } from '../../vscode/toast'
 
 export class PlotsModel extends ModelWithPersistence {
   private readonly experiments: Experiments
@@ -223,6 +227,23 @@ 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)
+
+    try {
+      writeJson(filePath, rawData as unknown as Record<string, unknown>, true)
+      void openFileInEditor(filePath)
+    } catch {
+      void Toast.showError('Cannot write to file')
+    }
+  }
+
   public getTemplatePlots(
     order: TemplateOrder | undefined,
     selectedRevisions: Revision[]
@@ -438,4 +459,24 @@ export class PlotsModel extends ModelWithPersistence {
       this.multiSourceEncoding
     )
   }
+
+  private getSelectedTemplatePlotData(path: string) {
+    const selectedRevisions = this.getSelectedRevisionDetails()
+
+    return collectSelectedTemplatePlotRawData({
+      multiSourceEncodingUpdate: this.multiSourceEncoding[path] || {},
+      path,
+      revisionData: this.revisionData,
+      selectedRevisions: selectedRevisions.map(({ id }) => id),
+      template: this.templates[path]
+    })
+  }
+
+  private getCustomPlotData(orderValue: CustomPlotsOrderValue) {
+    const experiments = this.experiments
+      .getWorkspaceCommitsAndExperiments()
+      .filter(({ id }) => id !== EXPERIMENT_WORKSPACE_ID)
+
+    return collectCustomPlotRawData(orderValue, experiments)
+  }
 }

extension/src/plots/webview/messages.ts

@@ -1,4 +1,4 @@
-import { commands } from 'vscode'
+import { Uri, commands } from 'vscode'
 import isEmpty from 'lodash.isempty'
 import {
   ComparisonPlot,
@@ -21,7 +21,11 @@ import {
 import { PlotsModel } from '../model'
 import { PathsModel } from '../paths/model'
 import { BaseWebview } from '../../webview'
-import { getModifiedTime, openImageFileInEditor } from '../../fileSystem'
+import {
+  getModifiedTime,
+  openImageFileInEditor,
+  showSaveDialog
+} from '../../fileSystem'
 import { reorderObjectList } from '../../util/array'
 import { CustomPlotsOrderValue } from '../model/custom'
 import { getCustomPlotId } from '../model/collect'
@@ -78,6 +82,8 @@ export class WebviewMessages {
           RegisteredCommands.PLOTS_CUSTOM_ADD,
           this.dvcRoot
         )
+      case MessageFromWebviewType.EXPORT_PLOT_DATA:
+        return this.exportPlotData(message.payload)
       case MessageFromWebviewType.RESIZE_PLOTS:
         return this.setPlotSize(
           message.payload.section,
@@ -337,4 +343,20 @@ export class WebviewMessages {
   private getCustomPlots() {
     return this.plots.getCustomPlots() || null
   }
+
+  private async exportPlotData(plotId: string) {
+    const file = await showSaveDialog(Uri.file('data.json'), { JSON: ['json'] })
+
+    if (!file) {
+      return
+    }
+
+    sendTelemetryEvent(
+      EventName.VIEWS_PLOTS_EXPORT_PLOT_DATA,
+      undefined,
+      undefined
+    )
+
+    this.plots.savePlotData(plotId, file.path)
+  }
 }

extension/src/telemetry/constants.ts

@@ -71,6 +71,7 @@ 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_FOCUS_CHANGED: 'views.plots.focusChanged',
     VIEWS_PLOTS_REVISIONS_REORDERED: 'views.plots.revisionsReordered',
     VIEWS_PLOTS_SECTION_RESIZED: 'views.plots.sectionResized',
@@ -264,6 +265,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_ZOOM_PLOT]: { isImage: boolean }
   [EventName.VIEWS_REORDER_PLOTS_CUSTOM]: undefined
   [EventName.VIEWS_REORDER_PLOTS_TEMPLATES]: undefined