Merge create plot commands (#4680)

extension/package.json

@@ -425,7 +425,7 @@
       },
       {
         "title": "Add Plot",
-        "command": "dvc.addTopLevelPlot",
+        "command": "dvc.addPlot",
         "category": "DVC",
         "icon": "$(add)"
       },
@@ -601,11 +601,6 @@
         "category": "DVC",
         "icon": "$(refresh)"
       },
-      {
-        "title": "Add Custom Plot",
-        "command": "dvc.views.plots.addCustomPlot",
-        "category": "DVC"
-      },
       {
         "title": "Remove Custom Plot(s)",
         "command": "dvc.views.plots.removeCustomPlots",
@@ -910,7 +905,7 @@
           "when": "dvc.commands.available && dvc.project.available"
         },
         {
-          "command": "dvc.addTopLevelPlot",
+          "command": "dvc.addPlot",
           "when": "dvc.commands.available && dvc.project.available"
         },
         {
@@ -1029,10 +1024,6 @@
           "command": "dvc.views.plots.refreshPlots",
           "when": "dvc.commands.available && dvc.project.available"
         },
-        {
-          "command": "dvc.views.plots.addCustomPlot",
-          "when": "dvc.commands.available && dvc.project.available"
-        },
         {
           "command": "dvc.views.plots.removeCustomPlots",
           "when": "dvc.commands.available && dvc.project.available"
@@ -1425,7 +1416,7 @@
           "group": "navigation@3"
         },
         {
-          "command": "dvc.addTopLevelPlot",
+          "command": "dvc.addPlot",
           "when": "view == dvc.views.plotsPathsTree",
           "group": "navigation@0"
         },

extension/src/commands/external.ts

@@ -72,15 +72,15 @@ 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',
   PLOTS_SELECT = 'dvc.views.plotsPathsTree.selectPlots',
   PLOTS_REFRESH = 'dvc.views.plots.refreshPlots',
-  PLOTS_CUSTOM_ADD = 'dvc.views.plots.addCustomPlot',
   PLOTS_CUSTOM_REMOVE = 'dvc.views.plots.removeCustomPlots',
 
+  ADD_PLOT = 'dvc.addPlot',
+
   EXPERIMENT_AND_PLOTS_SHOW = 'dvc.showExperimentsAndPlots',
 
   EXTENSION_CHECK_CLI_COMPATIBLE = 'dvc.checkCLICompatible',

extension/src/extension.ts

@@ -192,7 +192,12 @@ class Extension extends Disposable {
       this.setup
     )
     registerPipelineCommands(this.pipelines, this.internalCommands)
-    registerPlotsCommands(this.plots, this.internalCommands, this.setup)
+    registerPlotsCommands(
+      this.plots,
+      this.internalCommands,
+      this.setup,
+      this.pipelines
+    )
     registerSetupCommands(this.setup, this.internalCommands)
     this.internalCommands.registerExternalCommand(
       RegisteredCommands.EXPERIMENT_AND_PLOTS_SHOW,

extension/src/pipeline/index.ts

@@ -112,7 +112,7 @@ export class Pipeline extends DeferredDisposable {
     return this.addPipeline()
   }
 
-  public async addTopLevelPlot() {
+  public async addDataSeriesPlot() {
     const cwd = await this.getCwd()
 
     if (!cwd) {

extension/src/pipeline/register.ts

@@ -1,7 +1,6 @@
 import { WorkspacePipeline } from './workspace'
 import { RegisteredCommands } from '../commands/external'
 import { InternalCommands } from '../commands/internal'
-import { Context, getDvcRootFromContext } from '../vscode/context'
 
 export const registerPipelineCommands = (
   pipelines: WorkspacePipeline,
@@ -11,10 +10,4 @@ export const registerPipelineCommands = (
     RegisteredCommands.PIPELINE_SHOW_DAG,
     () => pipelines.showDag()
   )
-
-  internalCommands.registerExternalCommand(
-    RegisteredCommands.PIPELINE_ADD_PLOT,
-    (context: Context) =>
-      pipelines.addTopLevelPlot(getDvcRootFromContext(context))
-  )
 }

extension/src/pipeline/workspace.ts

@@ -63,14 +63,14 @@ export class WorkspacePipeline extends BaseWorkspace<Pipeline> {
     )
   }
 
-  public async addTopLevelPlot(overrideRoot?: string) {
+  public async addDataSeriesPlot(overrideRoot?: string) {
     const cwd = overrideRoot || (await this.getCwd())
 
     if (!cwd) {
       return
     }
 
-    void this.getRepository(cwd).addTopLevelPlot()
+    void this.getRepository(cwd).addDataSeriesPlot()
   }
 
   private getCwd() {

extension/src/plots/commands/register.ts

@@ -1,15 +1,22 @@
+import { commands } from 'vscode'
 import { RegisteredCommands } from '../../commands/external'
 import { InternalCommands } from '../../commands/internal'
 import { showSetupOrExecuteCommand } from '../../commands/util'
 import { Setup } from '../../setup'
 import { Context, getDvcRootFromContext } from '../../vscode/context'
 import { WorkspacePlots } from '../workspace'
+import { WorkspacePipeline } from '../../pipeline/workspace'
 
 export const registerPlotsCommands = (
   plots: WorkspacePlots,
   internalCommands: InternalCommands,
-  setup: Setup
+  setup: Setup,
+  pipelines: WorkspacePipeline
 ) => {
+  commands.registerCommand(RegisteredCommands.ADD_PLOT, (context: Context) =>
+    plots.addPlot(pipelines, getDvcRootFromContext(context))
+  )
+
   internalCommands.registerExternalCommand(
     RegisteredCommands.PLOTS_SHOW,
     showSetupOrExecuteCommand(setup, context =>
@@ -27,11 +34,6 @@ export const registerPlotsCommands = (
     (context: Context) => plots.refresh(getDvcRootFromContext(context))
   )
 
-  internalCommands.registerExternalCommand(
-    RegisteredCommands.PLOTS_CUSTOM_ADD,
-    (context: Context) => plots.addCustomPlot(getDvcRootFromContext(context))
-  )
-
   internalCommands.registerExternalCommand(
     RegisteredCommands.PLOTS_CUSTOM_REMOVE,
     (context: Context) =>

extension/src/plots/index.ts

@@ -223,7 +223,8 @@ export class Plots extends BaseRepository<TPlotsData> {
       errors,
       experiments,
       () => this.getWebview(),
-      () => this.selectPlots()
+      () => this.selectPlots(),
+      () => this.addCustomPlot()
     )
     this.dispose.track(
       this.onDidReceivedWebviewMessage(message =>

extension/src/plots/quickPick.test.ts

@@ -0,0 +1,39 @@
+import { pickPlotType } from './quickPick'
+import { quickPickValue } from '../vscode/quickPick'
+import { Title } from '../vscode/title'
+
+jest.mock('../vscode/quickPick')
+
+const mockedQuickPickValue = jest.mocked(quickPickValue)
+
+beforeEach(() => {
+  jest.resetAllMocks()
+})
+
+describe('pickPlotType', () => {
+  it('should call a quick pick with possible plot types', async () => {
+    mockedQuickPickValue.mockResolvedValueOnce(undefined)
+
+    await pickPlotType()
+
+    expect(mockedQuickPickValue).toHaveBeenCalledWith(
+      [
+        {
+          description:
+            'Create a data series plot definition by selecting data from a file',
+          label: 'Data Series',
+          value: 'data-series'
+        },
+        {
+          description:
+            'Create an extension-only plot by selecting one metric and one param from the experiments table',
+          label: 'Custom',
+          value: 'custom'
+        }
+      ],
+      {
+        title: Title.SELECT_PLOT_TYPE
+      }
+    )
+  })
+})

extension/src/plots/quickPick.ts

@@ -0,0 +1,28 @@
+import { quickPickValue } from '../vscode/quickPick'
+import { Title } from '../vscode/title'
+
+export const enum PLOT_TYPE {
+  CUSTOM = 'custom',
+  DATA_SERIES = 'data-series'
+}
+
+export const pickPlotType = () =>
+  quickPickValue(
+    [
+      {
+        description:
+          'Create a data series plot definition by selecting data from a file',
+        label: 'Data Series',
+        value: PLOT_TYPE.DATA_SERIES
+      },
+      {
+        description:
+          'Create an extension-only plot by selecting one metric and one param from the experiments table',
+        label: 'Custom',
+        value: PLOT_TYPE.CUSTOM
+      }
+    ],
+    {
+      title: Title.SELECT_PLOT_TYPE
+    }
+  )

extension/src/plots/webview/messages.ts

@@ -41,6 +41,7 @@ export class WebviewMessages {
 
   private readonly getWebview: () => BaseWebview<TPlotsData> | undefined
   private readonly selectPlots: () => Promise<void>
+  private readonly addCustomPlot: () => Promise<void>
 
   constructor(
     dvcRoot: string,
@@ -49,7 +50,8 @@ export class WebviewMessages {
     errors: ErrorsModel,
     experiments: Experiments,
     getWebview: () => BaseWebview<TPlotsData> | undefined,
-    selectPlots: () => Promise<void>
+    selectPlots: () => Promise<void>,
+    addCustomPlot: () => Promise<void>
   ) {
     this.dvcRoot = dvcRoot
     this.paths = paths
@@ -58,6 +60,7 @@ export class WebviewMessages {
     this.experiments = experiments
     this.getWebview = getWebview
     this.selectPlots = selectPlots
+    this.addCustomPlot = addCustomPlot
   }
 
   public async sendWebviewMessage() {
@@ -71,16 +74,13 @@ export class WebviewMessages {
 
   public handleMessageFromWebview(message: MessageFromWebview) {
     switch (message.type) {
-      case MessageFromWebviewType.ADD_CUSTOM_PLOT:
-        return commands.executeCommand(
-          RegisteredCommands.PLOTS_CUSTOM_ADD,
-          this.dvcRoot
-        )
-      case MessageFromWebviewType.ADD_PIPELINE_PLOT:
+      case MessageFromWebviewType.ADD_PLOT:
         return commands.executeCommand(
-          RegisteredCommands.PIPELINE_ADD_PLOT,
+          RegisteredCommands.ADD_PLOT,
           this.dvcRoot
         )
+      case MessageFromWebviewType.ADD_CUSTOM_PLOT:
+        return this.addCustomPlotFromWebview()
       case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_CSV:
         return this.exportPlotDataAsCsv(message.payload)
       case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_TSV:
@@ -309,6 +309,11 @@ export class WebviewMessages {
     )
   }
 
+  private addCustomPlotFromWebview() {
+    void this.addCustomPlot()
+    sendTelemetryEvent(EventName.VIEWS_PLOTS_CUSTOM_ADD, undefined, undefined)
+  }
+
   private setExperimentStatus(id: string) {
     this.experiments.toggleExperimentStatus(id)
     sendTelemetryEvent(

extension/src/plots/workspace.ts

@@ -1,9 +1,14 @@
 import { EventEmitter } from 'vscode'
 import { Plots } from '.'
 import { PlotsData } from './webview/contract'
+import { PLOT_TYPE, pickPlotType } from './quickPick'
 import { ResourceLocator } from '../resourceLocator'
 import { BaseWorkspaceWebviews } from '../webview/workspace'
 import { WorkspaceExperiments } from '../experiments/workspace'
+import { WorkspacePipeline } from '../pipeline/workspace'
+import { sendTelemetryEvent } from '../telemetry'
+import { RegisteredCommands } from '../commands/external'
+import { StopWatch } from '../util/time'
 
 export class WorkspacePlots extends BaseWorkspaceWebviews<Plots, PlotsData> {
   public readonly pathsChanged = this.dispose.track(new EventEmitter<void>())
@@ -74,6 +79,25 @@ export class WorkspacePlots extends BaseWorkspaceWebviews<Plots, PlotsData> {
     return this.getRepository(dvcRoot).removeCustomPlot()
   }
 
+  public async addPlot(pipelines: WorkspacePipeline, overrideRoot?: string) {
+    const stopWatch = new StopWatch()
+    const response: PLOT_TYPE | undefined = await pickPlotType()
+
+    if (response === PLOT_TYPE.CUSTOM) {
+      await this.addCustomPlot(overrideRoot)
+    }
+
+    if (response === PLOT_TYPE.DATA_SERIES) {
+      await pipelines.addDataSeriesPlot(overrideRoot)
+    }
+
+    sendTelemetryEvent(
+      RegisteredCommands.ADD_PLOT,
+      { type: response },
+      { duration: stopWatch.getElapsedTime() }
+    )
+  }
+
   public getFocusedOrOnlyOrPickProject() {
     return this.focusedWebviewDvcRoot || this.getOnlyOrPickProject()
   }

extension/src/telemetry/constants.ts

@@ -86,6 +86,7 @@ export const EventName = Object.assign(
     VIEWS_PLOTS_COMPARISON_ROWS_REORDERED:
       'views.plots.comparisonRowsReordered',
     VIEWS_PLOTS_CREATED: 'views.plots.created',
+    VIEWS_PLOTS_CUSTOM_ADD: 'views.plots.addCustomPlot',
     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',
@@ -195,15 +196,15 @@ export interface IEventNamePropertyMapping {
   [EventName.STOP_EXPERIMENTS]: { stopped: boolean; wasRunning: boolean }
 
   [EventName.PIPELINE_SHOW_DAG]: undefined
-  [EventName.PIPELINE_ADD_PLOT]: undefined
 
   [EventName.PLOTS_PATH_TOGGLE]: undefined
   [EventName.PLOTS_SHOW]: undefined
   [EventName.PLOTS_SELECT]: undefined
   [EventName.PLOTS_REFRESH]: undefined
-  [EventName.PLOTS_CUSTOM_ADD]: undefined
   [EventName.PLOTS_CUSTOM_REMOVE]: undefined
 
+  [EventName.ADD_PLOT]: { type: string | undefined }
+
   [EventName.ADD_TARGET]: undefined
   [EventName.CHECKOUT_TARGET]: undefined
   [EventName.CHECKOUT]: undefined
@@ -287,6 +288,7 @@ export interface IEventNamePropertyMapping {
 
   [EventName.VIEWS_PLOTS_CLOSED]: undefined
   [EventName.VIEWS_PLOTS_CREATED]: undefined
+  [EventName.VIEWS_PLOTS_CUSTOM_ADD]: undefined
   [EventName.VIEWS_PLOTS_FOCUS_CHANGED]: WebviewFocusChangedProperties
   [EventName.VIEWS_PLOTS_REVISIONS_REORDERED]: undefined
   [EventName.VIEWS_PLOTS_COMPARISON_ROWS_REORDERED]: undefined

extension/src/test/e2e/util.ts

@@ -185,7 +185,9 @@ export const closeAllEditors = async (): Promise<void> => {
 
 export const createCustomPlot = async (): Promise<void> => {
   const workbench = await browser.getWorkbench()
-  const addCustomPlot = await workbench.executeCommand('DVC: Add Custom Plot')
+  const addCustomPlot = await workbench.executeCommand('DVC: Add Plot')
+  await browser.waitUntil(() => addCustomPlot.elem.isDisplayed())
+  await browser.keys(['ArrowDown', 'Enter'])
   await browser.waitUntil(() => addCustomPlot.elem.isDisplayed())
   await browser.keys('Enter')
   await browser.waitUntil(() => addCustomPlot.elem.isDisplayed())