Merge create plot commands (#4680)

Author: julieg18

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
+    }
+  )