Add plotting actions to experiments table context menu (#2388)

* add plotting actions to experiments table context menu

* fix view title command registration

* fix brackets so that tests run and do what they are meant to

* update context menu labels

Author: mattseddon

extension/src/experiments/commands/register.ts

@@ -295,7 +295,8 @@ const registerExperimentRunCommands = (
 
   internalCommands.registerExternalCommand(
     RegisteredCommands.EXPERIMENT_SHOW,
-    () => experiments.showWebview()
+    (context: { dvcRoot?: string } | undefined) =>
+      experiments.showWebview(context?.dvcRoot)
   )
 }
 

extension/src/experiments/webview/messages.ts

@@ -114,7 +114,7 @@ export class WebviewMessages {
         return this.focusSortsTree()
 
       case MessageFromWebviewType.OPEN_PLOTS_WEBVIEW:
-        return commands.executeCommand(RegisteredCommands.PLOTS_SHOW)
+        return this.showPlots()
 
       case MessageFromWebviewType.SHARE_EXPERIMENT_AS_BRANCH:
         return commands.executeCommand(
@@ -127,6 +127,15 @@ export class WebviewMessages {
           { dvcRoot: this.dvcRoot, id: message.payload }
         )
 
+      case MessageFromWebviewType.SET_EXPERIMENTS_FOR_PLOTS:
+        return this.setSelectedExperiments(message.payload)
+
+      case MessageFromWebviewType.SET_EXPERIMENTS_AND_OPEN_PLOTS:
+        return Promise.all([
+          this.setSelectedExperiments(message.payload),
+          this.showPlots()
+        ])
+
       default:
         Logger.error(`Unexpected message: ${JSON.stringify(message)}`)
     }
@@ -254,4 +263,27 @@ export class WebviewMessages {
       undefined
     )
   }
+
+  private setSelectedExperiments(ids: string[]) {
+    const experiments = this.experiments
+      .getCombinedList()
+      .filter(({ id }) => ids.includes(id))
+
+    this.experiments.setSelectionMode(false)
+    this.experiments.setSelected(experiments)
+
+    this.notifyChanged()
+
+    sendTelemetryEvent(
+      EventName.VIEWS_EXPERIMENTS_TABLE_SELECT_EXPERIMENTS_FOR_PLOTS,
+      { experimentCount: ids.length },
+      undefined
+    )
+  }
+
+  private showPlots() {
+    return commands.executeCommand(RegisteredCommands.PLOTS_SHOW, {
+      dvcRoot: this.dvcRoot
+    })
+  }
 }

extension/src/plots/commands/register.ts

@@ -8,9 +8,8 @@ export const registerPlotsCommands = (
 ) => {
   internalCommands.registerExternalCommand(
     RegisteredCommands.PLOTS_SHOW,
-    () => {
-      plots.showWebview()
-    }
+    (context: { dvcRoot?: string } | undefined) =>
+      plots.showWebview(context?.dvcRoot)
   )
 
   internalCommands.registerExternalCommand(

extension/src/telemetry/constants.ts

@@ -51,6 +51,8 @@ export const EventName = Object.assign(
       'views.experimentsTable.columnResized',
     VIEWS_EXPERIMENTS_TABLE_SELECT_COLUMNS:
       'views.experimentsTable.selectColumns',
+    VIEWS_EXPERIMENTS_TABLE_SELECT_EXPERIMENTS_FOR_PLOTS:
+      'views.experimentsTable.selectExperimentsForPlots',
     VIEWS_EXPERIMENTS_TABLE_SORT_COLUMN:
       'views.experimentsTable.columnSortAdded',
 
@@ -217,6 +219,9 @@ export interface IEventNamePropertyMapping {
     path: string
   }
   [EventName.VIEWS_EXPERIMENTS_TABLE_SELECT_COLUMNS]: undefined
+  [EventName.VIEWS_EXPERIMENTS_TABLE_SELECT_EXPERIMENTS_FOR_PLOTS]: {
+    experimentCount: number
+  }
   [EventName.VIEWS_EXPERIMENTS_TABLE_OPEN_PARAMS_FILE]: {
     path: string
   }

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

@@ -64,6 +64,7 @@ import { ExperimentFlag } from '../../../cli/dvc/constants'
 import { DvcExecutor } from '../../../cli/dvc/executor'
 import { shortenForLabel } from '../../../util/string'
 import { GitExecutor } from '../../../cli/git/executor'
+import { WorkspacePlots } from '../../../plots/workspace'
 
 suite('Experiments Test Suite', () => {
   const disposable = Disposable.fn()
@@ -917,8 +918,76 @@ suite('Experiments Test Suite', () => {
         areExperimentsStarred(experimentsToToggle),
         'experiments have been starred'
       ).to.be.true
-    })
-  }).timeout(WEBVIEW_TEST_TIMEOUT)
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should be able to handle a message to select experiments for plotting', async () => {
+      const { experiments, experimentsModel } = buildExperiments(disposable)
+      await experiments.isReady()
+
+      const webview = await experiments.showWebview()
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+      const mockExperimentIds = [
+        'exp-e7a67',
+        'd1343a87c6ee4a2e82d19525964d2fb2cb6756c9',
+        'test-branch'
+      ]
+
+      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
+
+      const tableChangePromise = experimentsUpdatedEvent(experiments)
+
+      mockMessageReceived.fire({
+        payload: mockExperimentIds,
+        type: MessageFromWebviewType.SET_EXPERIMENTS_FOR_PLOTS
+      })
+
+      await tableChangePromise
+
+      const selectExperimentIds = experimentsModel
+        .getSelectedRevisions()
+        .map(({ id }) => id)
+        .sort()
+      expect(selectExperimentIds).to.deep.equal(mockExperimentIds.sort())
+      expect(selectExperimentIds).to.deep.equal(mockExperimentIds)
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should be able to handle a message to compare experiments plots', async () => {
+      const { experiments, experimentsModel } = buildExperiments(disposable)
+      const mockShowPlots = stub(
+        WorkspacePlots.prototype,
+        'showWebview'
+      ).resolves(undefined)
+
+      await experiments.isReady()
+
+      const webview = await experiments.showWebview()
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+      const mockExperimentIds = [
+        'exp-e7a67',
+        'd1343a87c6ee4a2e82d19525964d2fb2cb6756c9',
+        'test-branch'
+      ]
+
+      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
+
+      const tableChangePromise = experimentsUpdatedEvent(experiments)
+
+      mockMessageReceived.fire({
+        payload: mockExperimentIds,
+        type: MessageFromWebviewType.SET_EXPERIMENTS_AND_OPEN_PLOTS
+      })
+
+      await tableChangePromise
+
+      const selectExperimentIds = experimentsModel
+        .getSelectedRevisions()
+        .map(({ id }) => id)
+        .sort()
+      expect(selectExperimentIds).to.deep.equal(mockExperimentIds.sort())
+      expect(mockShowPlots).to.be.calledOnce
+      expect(mockShowPlots).to.be.calledWith(dvcDemoPath)
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+  })
 
   describe('Sorting', () => {
     it('should be able to sort', async () => {

extension/src/webview/contract.ts

@@ -36,6 +36,8 @@ export enum MessageFromWebviewType {
   SELECT_EXPERIMENTS = 'select-experiments',
   SELECT_COLUMNS = 'select-columns',
   SELECT_PLOTS = 'select-plots',
+  SET_EXPERIMENTS_FOR_PLOTS = 'set-experiments-for-plots',
+  SET_EXPERIMENTS_AND_OPEN_PLOTS = 'set-experiments-and-open-plots',
   SHARE_EXPERIMENT_AS_BRANCH = 'share-experiment-as-branch',
   SHARE_EXPERIMENT_AS_COMMIT = 'share-experiment-as-commit',
   TOGGLE_METRIC = 'toggle-metric',
@@ -153,6 +155,14 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.FOCUS_FILTERS_TREE }
   | { type: MessageFromWebviewType.FOCUS_SORTS_TREE }
   | { type: MessageFromWebviewType.OPEN_PLOTS_WEBVIEW }
+  | {
+      type: MessageFromWebviewType.SET_EXPERIMENTS_FOR_PLOTS
+      payload: string[]
+    }
+  | {
+      type: MessageFromWebviewType.SET_EXPERIMENTS_AND_OPEN_PLOTS
+      payload: string[]
+    }
   | {
       type: MessageFromWebviewType.SHARE_EXPERIMENT_AS_BRANCH
       payload: string

extension/src/webview/workspace.ts

@@ -27,8 +27,8 @@ export abstract class BaseWorkspaceWebviews<
     }
   }
 
-  public async showWebview() {
-    const dvcRoot = await this.getOnlyOrPickProject()
+  public async showWebview(overrideRoot?: string) {
+    const dvcRoot = overrideRoot || (await this.getOnlyOrPickProject())
     if (!dvcRoot) {
       return
     }

webview/src/experiments/components/App.test.tsx

@@ -868,6 +868,22 @@ describe('App', () => {
       expect(itemLabels).toContain('Remove Selected Rows')
     })
 
+    it('should always present the Plots options if multiple rows are selected', () => {
+      renderTableWithoutRunningExperiments()
+
+      clickRowCheckbox('4fb124a')
+      clickRowCheckbox('42b8736')
+
+      const target = screen.getByText('4fb124a')
+      fireEvent.contextMenu(target, { bubbles: true })
+
+      jest.advanceTimersByTime(100)
+      const menuitems = screen.getAllByRole('menuitem')
+      const itemLabels = menuitems.map(item => item.textContent)
+      expect(itemLabels).toContain('Plot and Show')
+      expect(itemLabels).toContain('Plot')
+    })
+
     it('should allow batch selection of rows by shift-clicking a range of them', () => {
       renderTableWithoutRunningExperiments()
 

webview/src/experiments/components/table/RowContextMenu.tsx

@@ -7,7 +7,7 @@ import { MessagesMenuOptionProps } from '../../../shared/components/messagesMenu
 import { cond } from '../../../util/helpers'
 
 const experimentMenuOption = (
-  id: string | string[],
+  payload: string | string[],
   label: string,
   type: MessageFromWebviewType,
   hidden?: boolean,
@@ -19,7 +19,7 @@ const experimentMenuOption = (
     id: type,
     label,
     message: {
-      payload: id,
+      payload,
       type
     }
   } as MessagesMenuOptionProps
@@ -45,6 +45,8 @@ const getMultiSelectMenuOptions = (
     }) => starred
   )
 
+  const selectedIds = selectedRowsList.map(value => value.row.values.id)
+
   const removableRowIds = selectedRowsList
     .filter(value => value.row.depth === 1)
     .map(value => value.row.values.id)
@@ -69,6 +71,20 @@ const getMultiSelectMenuOptions = (
       starredExperiments.map(value => value.row.values.id),
       'Unstar'
     ),
+    experimentMenuOption(
+      selectedIds,
+      'Plot',
+      MessageFromWebviewType.SET_EXPERIMENTS_FOR_PLOTS,
+      false,
+      true
+    ),
+    experimentMenuOption(
+      selectedIds,
+      'Plot and Show',
+      MessageFromWebviewType.SET_EXPERIMENTS_AND_OPEN_PLOTS,
+      false,
+      false
+    ),
     experimentMenuOption(
       removableRowIds,
       'Remove Selected Rows',