Add custom plot commands to the command pallete (#3629)

Author: julieg18

extension/package.json

@@ -564,6 +564,16 @@
         "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",
+        "category": "DVC"
+      },
       {
         "title": "Reset Persisted State and Reload Window",
         "command": "dvc.resetState",
@@ -947,6 +957,14 @@
         {
           "command": "dvc.views.plotsPathsTree.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"
         }
       ],
       "scm/title": [

extension/src/commands/external.ts

@@ -71,6 +71,8 @@ export enum RegisteredCommands {
   PLOTS_SHOW = 'dvc.showPlots',
   PLOTS_SELECT = 'dvc.views.plotsPathsTree.selectPlots',
   PLOTS_REFRESH = 'dvc.views.plotsPathsTree.refreshPlots',
+  PLOTS_CUSTOM_ADD = 'dvc.views.plots.addCustomPlot',
+  PLOTS_CUSTOM_REMOVE = 'dvc.views.plots.removeCustomPlots',
 
   EXPERIMENT_AND_PLOTS_SHOW = 'dvc.showExperimentsAndPlots',
 

extension/src/plots/commands/register.ts

@@ -26,4 +26,15 @@ export const registerPlotsCommands = (
     RegisteredCommands.PLOTS_REFRESH,
     (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) =>
+      plots.removeCustomPlots(getDvcRootFromContext(context))
+  )
 }

extension/src/plots/index.ts

@@ -7,6 +7,7 @@ import { ErrorsModel } from './errors/model'
 import { PlotsModel } from './model'
 import { collectEncodingElements, collectScale } from './paths/collect'
 import { PathsModel } from './paths/model'
+import { pickCustomPlots, pickMetricAndParam } from './model/quickPick'
 import { BaseWebview } from '../webview'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
@@ -19,6 +20,7 @@ import { Toast } from '../vscode/toast'
 import { pickPaths } from '../path/selection/quickPick'
 import { ErrorDecorationProvider } from '../tree/decorationProvider/error'
 import { DecoratableTreeItemScheme } from '../tree'
+import { Title } from '../vscode/title'
 
 export type PlotsWebview = BaseWebview<TPlotsData>
 
@@ -29,6 +31,7 @@ export class Plots extends BaseRepository<TPlotsData> {
 
   private readonly pathsChanged = this.dispose.track(new EventEmitter<void>())
 
+  private readonly experiments: Experiments
   private readonly plots: PlotsModel
   private readonly paths: PathsModel
   private readonly data: PlotsData
@@ -56,6 +59,7 @@ export class Plots extends BaseRepository<TPlotsData> {
     this.paths = this.dispose.track(
       new PathsModel(this.dvcRoot, this.errors, workspaceState)
     )
+    this.experiments = experiments
 
     this.webviewMessages = this.createWebviewMessageHandler(
       this.paths,
@@ -112,6 +116,37 @@ export class Plots extends BaseRepository<TPlotsData> {
     this.triggerDataUpdate()
   }
 
+  public async addCustomPlot() {
+    const metricAndParam = await pickMetricAndParam(
+      this.experiments.getColumnTerminalNodes(),
+      this.plots.getCustomPlotsOrder()
+    )
+
+    if (!metricAndParam) {
+      return
+    }
+
+    this.plots.addCustomPlot(metricAndParam)
+    void this.sendPlots()
+  }
+
+  public async removeCustomPlot() {
+    const selectedPlotsIds = await pickCustomPlots(
+      this.plots.getCustomPlotsOrder(),
+      'There are no plots to remove.',
+      {
+        title: Title.SELECT_CUSTOM_PLOTS_TO_REMOVE
+      }
+    )
+
+    if (!selectedPlotsIds) {
+      return
+    }
+
+    this.plots.removeCustomPlots(selectedPlotsIds)
+    void this.sendPlots()
+  }
+
   public getChildPaths(path: string | undefined) {
     const multiSourceEncoding = this.plots.getMultiSourceData()
 

extension/src/plots/webview/messages.ts

@@ -1,3 +1,4 @@
+import { commands } from 'vscode'
 import isEmpty from 'lodash.isempty'
 import {
   ComparisonPlot,
@@ -21,12 +22,11 @@ import {
 import { PlotsModel } from '../model'
 import { PathsModel } from '../paths/model'
 import { BaseWebview } from '../../webview'
-import { pickCustomPlots, pickMetricAndParam } from '../model/quickPick'
 import { getModifiedTime, openImageFileInEditor } from '../../fileSystem'
-import { Title } from '../../vscode/title'
 import { reorderObjectList } from '../../util/array'
 import { CustomPlotsOrderValue } from '../model/custom'
 import { getCustomPlotId } from '../model/collect'
+import { RegisteredCommands } from '../../commands/external'
 import { ErrorsModel } from '../errors/model'
 
 export class WebviewMessages {
@@ -74,7 +74,7 @@ export class WebviewMessages {
   public handleMessageFromWebview(message: MessageFromWebview) {
     switch (message.type) {
       case MessageFromWebviewType.ADD_CUSTOM_PLOT:
-        return this.addCustomPlot()
+        return commands.executeCommand(RegisteredCommands.PLOTS_CUSTOM_ADD)
       case MessageFromWebviewType.RESIZE_PLOTS:
         return this.setPlotSize(
           message.payload.section,
@@ -96,7 +96,7 @@ export class WebviewMessages {
       case MessageFromWebviewType.SELECT_EXPERIMENTS:
         return this.selectExperimentsFromWebview()
       case MessageFromWebviewType.REMOVE_CUSTOM_PLOTS:
-        return this.removeCustomPlots()
+        return commands.executeCommand(RegisteredCommands.PLOTS_CUSTOM_REMOVE)
       case MessageFromWebviewType.REFRESH_REVISIONS:
         return this.refreshData()
       case MessageFromWebviewType.TOGGLE_EXPERIMENT:
@@ -183,37 +183,6 @@ export class WebviewMessages {
     )
   }
 
-  private async addCustomPlot() {
-    const metricAndParam = await pickMetricAndParam(
-      this.experiments.getColumnTerminalNodes(),
-      this.plots.getCustomPlotsOrder()
-    )
-
-    if (!metricAndParam) {
-      return
-    }
-
-    this.plots.addCustomPlot(metricAndParam)
-    this.sendCustomPlotsAndEvent(EventName.VIEWS_PLOTS_CUSTOM_PLOT_ADDED)
-  }
-
-  private async removeCustomPlots() {
-    const selectedPlotsIds = await pickCustomPlots(
-      this.plots.getCustomPlotsOrder(),
-      'There are no plots to remove.',
-      {
-        title: Title.SELECT_CUSTOM_PLOTS_TO_REMOVE
-      }
-    )
-
-    if (!selectedPlotsIds) {
-      return
-    }
-
-    this.plots.removeCustomPlots(selectedPlotsIds)
-    this.sendCustomPlotsAndEvent(EventName.VIEWS_PLOTS_CUSTOM_PLOT_REMOVED)
-  }
-
   private setCustomPlotsOrder(plotIds: string[]) {
     const customPlotsOrderWithId = this.plots
       .getCustomPlotsOrder()
@@ -232,17 +201,12 @@ export class WebviewMessages {
     }))
 
     this.plots.setCustomPlotsOrder(newOrder)
-    this.sendCustomPlotsAndEvent(EventName.VIEWS_REORDER_PLOTS_CUSTOM)
-  }
-
-  private sendCustomPlotsAndEvent(
-    event:
-      | typeof EventName.VIEWS_PLOTS_CUSTOM_PLOT_ADDED
-      | typeof EventName.VIEWS_PLOTS_CUSTOM_PLOT_REMOVED
-      | typeof EventName.VIEWS_REORDER_PLOTS_CUSTOM
-  ) {
     this.sendCustomPlots()
-    sendTelemetryEvent(event, undefined, undefined)
+    sendTelemetryEvent(
+      EventName.VIEWS_REORDER_PLOTS_CUSTOM,
+      undefined,
+      undefined
+    )
   }
 
   private selectPlotsFromWebview() {

extension/src/plots/workspace.ts

@@ -58,6 +58,22 @@ export class WorkspacePlots extends BaseWorkspaceWebviews<Plots, PlotsData> {
     return this.getRepository(dvcRoot).selectPlots()
   }
 
+  public async addCustomPlot(overrideRoot?: string) {
+    const dvcRoot = await this.getDvcRoot(overrideRoot)
+    if (!dvcRoot) {
+      return
+    }
+    return this.getRepository(dvcRoot).addCustomPlot()
+  }
+
+  public async removeCustomPlots(overrideRoot?: string) {
+    const dvcRoot = await this.getDvcRoot(overrideRoot)
+    if (!dvcRoot) {
+      return
+    }
+    return this.getRepository(dvcRoot).removeCustomPlot()
+  }
+
   public getFocusedOrOnlyOrPickProject() {
     return this.focusedWebviewDvcRoot || this.getOnlyOrPickProject()
   }

extension/src/telemetry/constants.ts

@@ -63,8 +63,6 @@ export const EventName = Object.assign(
     VIEWS_PLOTS_COMPARISON_ROWS_REORDERED:
       'views.plots.comparisonRowsReordered',
     VIEWS_PLOTS_CREATED: 'views.plots.created',
-    VIEWS_PLOTS_CUSTOM_PLOT_ADDED: 'views.plots.addCustomPlot',
-    VIEWS_PLOTS_CUSTOM_PLOT_REMOVED: 'views.plots.removeCustomPlot',
     VIEWS_PLOTS_EXPERIMENT_TOGGLE: 'views.plots.toggleExperimentStatus',
     VIEWS_PLOTS_FOCUS_CHANGED: 'views.plots.focusChanged',
     VIEWS_PLOTS_MANUAL_REFRESH: 'views.plots.manualRefresh',
@@ -172,6 +170,8 @@ export interface IEventNamePropertyMapping {
   [EventName.PLOTS_SHOW]: undefined
   [EventName.PLOTS_SELECT]: undefined
   [EventName.PLOTS_REFRESH]: undefined
+  [EventName.PLOTS_CUSTOM_ADD]: undefined
+  [EventName.PLOTS_CUSTOM_REMOVE]: undefined
 
   [EventName.ADD_TARGET]: undefined
   [EventName.CHECKOUT_TARGET]: undefined
@@ -240,8 +240,6 @@ export interface IEventNamePropertyMapping {
 
   [EventName.VIEWS_PLOTS_CLOSED]: undefined
   [EventName.VIEWS_PLOTS_CREATED]: undefined
-  [EventName.VIEWS_PLOTS_CUSTOM_PLOT_REMOVED]: undefined
-  [EventName.VIEWS_PLOTS_CUSTOM_PLOT_ADDED]: undefined
   [EventName.VIEWS_PLOTS_FOCUS_CHANGED]: WebviewFocusChangedProperties
   [EventName.VIEWS_PLOTS_MANUAL_REFRESH]: undefined
   [EventName.VIEWS_PLOTS_REVISIONS_REORDERED]: undefined