Add multi-select versions of exp push (#3792)

Author: mattseddon

extension/package.json

@@ -511,6 +511,12 @@
         "category": "DVC",
         "icon": "$(repo-push)"
       },
+      {
+        "title": "Push",
+        "command": "dvc.views.experimentsTree.pushExperiment",
+        "category": "DVC",
+        "icon": "$(repo-push)"
+      },
       {
         "title": "Show Logs",
         "command": "dvc.views.experiments.showLogs",
@@ -863,6 +869,14 @@
           "command": "dvc.views.experiments.branchExperiment",
           "when": "false"
         },
+        {
+          "command": "dvc.views.experiments.pushExperiment",
+          "when": "false"
+        },
+        {
+          "command": "dvc.views.experimentsTree.pushExperiment",
+          "when": "false"
+        },
         {
           "command": "dvc.views.experiments.queueExperiment",
           "when": "false"
@@ -887,10 +901,6 @@
           "command": "dvc.views.experiments.resetAndRunCheckpointExperiment",
           "when": "false"
         },
-        {
-          "command": "dvc.views.experiments.pushExperiment",
-          "when": "false"
-        },
         {
           "command": "dvc.views.experiments.showLogs",
           "when": "false"
@@ -1152,7 +1162,7 @@
           "when": "view == dvc.views.experimentsTree && dvc.commands.available && viewItem =~ /^(experiment|queued)$/ && !dvc.experiment.running"
         },
         {
-          "command": "dvc.views.experiments.pushExperiment",
+          "command": "dvc.views.experimentsTree.pushExperiment",
           "group": "1_share@0",
           "when": "view == dvc.views.experimentsTree && dvc.commands.available && viewItem == experiment && !dvc.experiment.running"
         },

extension/src/experiments/commands/index.ts

@@ -44,7 +44,7 @@ const convertUrlTextToLink = (stdout: string) => {
 
 export const getPushExperimentCommand =
   (internalCommands: InternalCommands, setup: Setup) =>
-  ({ dvcRoot, id }: { dvcRoot: string; id: string }) => {
+  ({ dvcRoot, ids }: { dvcRoot: string; ids: string[] }) => {
     const studioAccessToken = setup.getStudioAccessToken()
     if (
       !(
@@ -58,14 +58,14 @@ export const getPushExperimentCommand =
     return Toast.showProgress('exp push', async progress => {
       progress.report({ increment: 0 })
 
-      progress.report({ increment: 25, message: `Pushing ${id}...` })
+      progress.report({ increment: 25, message: `Pushing ${ids.join(' ')}...` })
 
       const remainingProgress = 75
 
       const stdout = await internalCommands.executeCommand(
         AvailableCommands.EXP_PUSH,
         dvcRoot,
-        id
+        ...ids
       )
 
       progress.report({

extension/src/experiments/model/collect.ts

@@ -328,10 +328,10 @@ export const collectExperiments = (
   return acc
 }
 
-type DeletableExperimentAccumulator = { [dvcRoot: string]: Set<string> }
+type ExperimentTypesAccumulator = { [dvcRoot: string]: Set<string> }
 
 const initializeAccumulatorRoot = (
-  acc: DeletableExperimentAccumulator,
+  acc: ExperimentTypesAccumulator,
   dvcRoot: string
 ) => {
   if (!acc[dvcRoot]) {
@@ -340,12 +340,12 @@ const initializeAccumulatorRoot = (
 }
 
 const collectExperimentItem = (
-  acc: DeletableExperimentAccumulator,
-  deletable: Set<string>,
+  acc: ExperimentTypesAccumulator,
+  types: Set<string>,
   experimentItem: ExperimentItem
 ) => {
   const { dvcRoot, type, id, label } = experimentItem
-  if (!deletable.has(type)) {
+  if (!types.has(type)) {
     return
   }
   initializeAccumulatorRoot(acc, dvcRoot)
@@ -357,18 +357,17 @@ const collectExperimentItem = (
   acc[dvcRoot].add(id)
 }
 
-export const collectDeletable = (
-  experimentItems: (string | ExperimentItem)[]
-): DeletableExperimentAccumulator => {
-  const deletable = new Set([ExperimentType.EXPERIMENT, ExperimentType.QUEUED])
-
-  const acc: DeletableExperimentAccumulator = {}
+export const collectExperimentType = (
+  experimentItems: (string | ExperimentItem)[],
+  types: Set<ExperimentType>
+): ExperimentTypesAccumulator => {
+  const acc: ExperimentTypesAccumulator = {}
   for (const experimentItem of experimentItems) {
     if (typeof experimentItem === 'string') {
       continue
     }
 
-    collectExperimentItem(acc, deletable, experimentItem)
+    collectExperimentItem(acc, types, experimentItem)
   }
 
   return acc

extension/src/experiments/model/tree.ts

@@ -9,7 +9,7 @@ import {
   Uri
 } from 'vscode'
 import { ExperimentType } from '.'
-import { collectDeletable, ExperimentItem } from './collect'
+import { collectExperimentType, ExperimentItem } from './collect'
 import { MAX_SELECTED_EXPERIMENTS } from './status'
 import { getDataFromColumnPaths } from './util'
 import { WorkspaceExperiments } from '../workspace'
@@ -69,7 +69,7 @@ export class ExperimentsTree
     this.experiments = experiments
     this.resourceLocator = resourceLocator
 
-    this.registerWorkaroundCommand()
+    this.registerWorkaroundCommands()
 
     this.updateDescriptionOnChange()
   }
@@ -123,24 +123,44 @@ export class ExperimentsTree
     )
   }
 
-  private registerWorkaroundCommand() {
+  private registerWorkaroundCommands() {
+    const callCommandWithSelected = async (
+      command:
+        | RegisteredCliCommands.EXPERIMENT_VIEW_REMOVE
+        | RegisteredCliCommands.EXPERIMENT_VIEW_PUSH,
+      experimentItem: ExperimentItem | string,
+      types: ExperimentType[]
+    ) => {
+      const selected = [
+        ...this.getSelectedExperimentItems(),
+        experimentItem
+      ] as (string | ExperimentItem)[]
+
+      const acc = collectExperimentType(selected, new Set(types))
+
+      for (const [dvcRoot, ids] of Object.entries(acc)) {
+        await commands.executeCommand(command, { dvcRoot, ids: [...ids] })
+      }
+    }
+
     commands.registerCommand(
       'dvc.views.experimentsTree.removeExperiment',
-      async experimentItem => {
-        const selected = [
-          ...this.getSelectedExperimentItems(),
-          experimentItem
-        ] as (string | ExperimentItem)[]
-
-        const deletable = collectDeletable(selected)
-
-        for (const [dvcRoot, ids] of Object.entries(deletable)) {
-          await commands.executeCommand(
-            RegisteredCliCommands.EXPERIMENT_VIEW_REMOVE,
-            { dvcRoot, ids }
-          )
-        }
-      }
+      (experimentItem: ExperimentItem) =>
+        callCommandWithSelected(
+          RegisteredCliCommands.EXPERIMENT_VIEW_REMOVE,
+          experimentItem,
+          [ExperimentType.EXPERIMENT, ExperimentType.QUEUED]
+        )
+    )
+
+    commands.registerCommand(
+      'dvc.views.experimentsTree.pushExperiment',
+      (experimentItem: ExperimentItem) =>
+        callCommandWithSelected(
+          RegisteredCliCommands.EXPERIMENT_VIEW_PUSH,
+          experimentItem,
+          [ExperimentType.EXPERIMENT]
+        )
     )
   }
 

extension/src/experiments/webview/messages.ts

@@ -197,7 +197,7 @@ export class WebviewMessages {
       case MessageFromWebviewType.PUSH_EXPERIMENT:
         return commands.executeCommand(
           RegisteredCliCommands.EXPERIMENT_VIEW_PUSH,
-          { dvcRoot: this.dvcRoot, id: message.payload }
+          { dvcRoot: this.dvcRoot, ids: message.payload }
         )
 
       case MessageFromWebviewType.SHOW_EXPERIMENT_LOGS:

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

@@ -671,7 +671,7 @@ suite('Experiments Test Suite', () => {
       )
 
       mockMessageReceived.fire({
-        payload: mockExpId,
+        payload: [mockExpId],
         type: MessageFromWebviewType.PUSH_EXPERIMENT
       })
 
@@ -712,7 +712,7 @@ suite('Experiments Test Suite', () => {
       )
 
       mockMessageReceived.fire({
-        payload: mockExpId,
+        payload: [mockExpId],
         type: MessageFromWebviewType.PUSH_EXPERIMENT
       })
 

extension/src/test/suite/experiments/model/tree.test.ts

@@ -12,6 +12,7 @@ import { ExperimentType } from '../../../../experiments/model'
 import { UNSELECTED } from '../../../../experiments/model/status'
 import {
   bypassProcessManagerDebounce,
+  bypassProgressCloseDelay,
   getFirstArgOfLastCall,
   getMockNow,
   getTimeSafeDisposer,
@@ -345,6 +346,99 @@ suite('Experiments Tree Test Suite', () => {
       )
     })
 
+    it('should be able to push an experiment with dvc.views.experimentsTree.pushExperiment', async () => {
+      bypassProgressCloseDelay()
+      const mockExperimentId = 'exp-to-push'
+      const mockExperiment = {
+        dvcRoot: dvcDemoPath,
+        id: mockExperimentId,
+        type: ExperimentType.EXPERIMENT
+      }
+
+      const mockExpPush = stub(DvcExecutor.prototype, 'expPush').resolves('')
+
+      stubPrivatePrototypeMethod(
+        ExperimentsTree,
+        'getSelectedExperimentItems'
+      ).returns([mockExperiment])
+
+      await commands.executeCommand(
+        'dvc.views.experimentsTree.pushExperiment',
+        mockExperiment
+      )
+
+      expect(mockExpPush).to.be.calledWithExactly(dvcDemoPath, mockExperimentId)
+    })
+
+    it('should be able to push the provided experiment with dvc.views.experimentsTree.pushExperiment (if no experiments are selected)', async () => {
+      bypassProgressCloseDelay()
+      const mockExperiment = 'exp-to-push'
+
+      const mockExpPush = stub(DvcExecutor.prototype, 'expPush').resolves('')
+
+      stubPrivatePrototypeMethod(
+        ExperimentsTree,
+        'getSelectedExperimentItems'
+      ).returns([])
+
+      await commands.executeCommand(
+        'dvc.views.experimentsTree.pushExperiment',
+        {
+          dvcRoot: dvcDemoPath,
+          id: mockExperiment,
+          type: ExperimentType.EXPERIMENT
+        }
+      )
+
+      expect(mockExpPush).to.be.calledWithExactly(dvcDemoPath, mockExperiment)
+    })
+
+    it('should be able to push multiple experiments with dvc.views.experimentsTree.pushExperiment', async () => {
+      bypassProgressCloseDelay()
+      const mockFirstExperimentId = 'first-exp-pushed'
+      const mockSecondExperimentId = 'second-exp-pushed'
+      const mockQueuedExperimentLabel = 'queued-excluded'
+
+      const mockExpPush = stub(DvcExecutor.prototype, 'expPush').resolves('')
+
+      stubPrivatePrototypeMethod(
+        ExperimentsTree,
+        'getSelectedExperimentItems'
+      ).returns([
+        dvcDemoPath,
+        {
+          dvcRoot: dvcDemoPath,
+          label: mockQueuedExperimentLabel,
+          type: ExperimentType.QUEUED
+        },
+        {
+          dvcRoot: dvcDemoPath,
+          id: mockFirstExperimentId,
+          type: ExperimentType.EXPERIMENT
+        },
+        {
+          dvcRoot: dvcDemoPath,
+          id: 'workspace-excluded',
+          type: ExperimentType.WORKSPACE
+        }
+      ])
+
+      await commands.executeCommand(
+        'dvc.views.experimentsTree.pushExperiment',
+        {
+          dvcRoot: dvcDemoPath,
+          id: mockSecondExperimentId,
+          type: ExperimentType.EXPERIMENT
+        }
+      )
+
+      expect(mockExpPush).to.be.calledWithExactly(
+        dvcDemoPath,
+        mockFirstExperimentId,
+        mockSecondExperimentId
+      )
+    })
+
     it('should be able to apply an experiment to the workspace with dvc.views.experiments.applyExperiment', async () => {
       const { experiments } = buildExperiments(disposable)
 

extension/src/webview/contract.ts

@@ -155,7 +155,7 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.SHOW_EXPERIMENT_LOGS; payload: string }
   | {
       type: MessageFromWebviewType.PUSH_EXPERIMENT
-      payload: string
+      payload: string[]
     }
   | {
       type: MessageFromWebviewType.REMOVE_COLUMN_SORT