Add multi-select remove to experiments tree (#1810)

* add multi select remove to the experiments tree

* refactor and test further

* cover case where icon clicked does not belong to a selected experiment

* clean up test

Author: mattseddon

extension/src/cli/executor.ts

@@ -76,11 +76,11 @@ export class CliExecutor extends Cli {
     )
   }
 
-  public experimentRemove(cwd: string, experimentName: string) {
+  public experimentRemove(cwd: string, ...experimentNames: string[]) {
     return this.executeExperimentProcess(
       cwd,
       ExperimentSubCommand.REMOVE,
-      experimentName
+      ...experimentNames
     )
   }
 

extension/src/experiments/model/collect.ts

@@ -1,4 +1,6 @@
+import { ThemeIcon, TreeItemCollapsibleState, Uri } from 'vscode'
 import omit from 'lodash.omit'
+import { ExperimentType } from '.'
 import { ExperimentsAccumulator } from './accumulator'
 import { extractColumns } from '../columns/extract'
 import { Experiment, ColumnType } from '../webview/contract'
@@ -10,6 +12,23 @@ import {
 } from '../../cli/reader'
 import { addToMapArray } from '../../util/map'
 import { uniqueValues } from '../../util/array'
+import { RegisteredCommands } from '../../commands/external'
+import { Resource } from '../../resourceLocator'
+
+export type ExperimentItem = {
+  command?: {
+    arguments: { dvcRoot: string; id: string }[]
+    command: RegisteredCommands
+    title: string
+  }
+  dvcRoot: string
+  description: string | undefined
+  id: string
+  label: string
+  collapsibleState: TreeItemCollapsibleState
+  type: ExperimentType
+  iconPath: ThemeIcon | Uri | Resource
+}
 
 type ExperimentsObject = { [sha: string]: ExperimentFieldsOrError }
 
@@ -288,3 +307,49 @@ export const collectMutableRevisions = (
 
   return uniqueValues(acc)
 }
+
+type DeletableExperimentAccumulator = { [dvcRoot: string]: Set<string> }
+
+const initializeAccumulatorRoot = (
+  acc: DeletableExperimentAccumulator,
+  dvcRoot: string
+) => {
+  if (!acc[dvcRoot]) {
+    acc[dvcRoot] = new Set<string>()
+  }
+}
+
+const collectExperimentItem = (
+  acc: DeletableExperimentAccumulator,
+  deletable: Set<string>,
+  experimentItem: ExperimentItem
+) => {
+  const { dvcRoot, type, id, label } = experimentItem
+  if (!deletable.has(type)) {
+    return
+  }
+  initializeAccumulatorRoot(acc, dvcRoot)
+  if (type === ExperimentType.QUEUED) {
+    acc[dvcRoot].add(label)
+    return
+  }
+
+  acc[dvcRoot].add(id)
+}
+
+export const collectDeletable = (
+  experimentItems: (string | ExperimentItem)[]
+): DeletableExperimentAccumulator => {
+  const deletable = new Set([ExperimentType.EXPERIMENT, ExperimentType.QUEUED])
+
+  const acc: DeletableExperimentAccumulator = {}
+  for (const experimentItem of experimentItems) {
+    if (typeof experimentItem === 'string') {
+      continue
+    }
+
+    collectExperimentItem(acc, deletable, experimentItem)
+  }
+
+  return acc
+}

extension/src/experiments/model/tree.ts

@@ -8,6 +8,7 @@ import {
   Uri
 } from 'vscode'
 import { ExperimentType } from '.'
+import { collectDeletable, ExperimentItem } from './collect'
 import { MAX_SELECTED_EXPERIMENTS } from './status'
 import { WorkspaceExperiments } from '../workspace'
 import { sendViewOpenedTelemetryEvent } from '../../telemetry'
@@ -21,21 +22,6 @@ import { sum } from '../../util/math'
 import { Title } from '../../vscode/title'
 import { Disposable } from '../../class/dispose'
 
-export type ExperimentItem = {
-  command?: {
-    arguments: { dvcRoot: string; id: string }[]
-    command: RegisteredCommands
-    title: string
-  }
-  dvcRoot: string
-  description: string | undefined
-  id: string
-  label: string
-  collapsibleState: TreeItemCollapsibleState
-  type: ExperimentType
-  iconPath: ThemeIcon | Uri | Resource
-}
-
 export class ExperimentsTree
   extends Disposable
   implements TreeDataProvider<string | ExperimentItem>
@@ -60,7 +46,7 @@ export class ExperimentsTree
     this.onDidChangeTreeData = experiments.experimentsChanged.event
 
     this.view = this.dispose.track(
-      createTreeView<ExperimentItem>('dvc.views.experimentsTree', this)
+      createTreeView<ExperimentItem>('dvc.views.experimentsTree', this, true)
     )
 
     this.dispose.track(
@@ -183,12 +169,19 @@ export class ExperimentsTree
 
     internalCommands.registerExternalCommand<ExperimentItem>(
       RegisteredCommands.EXPERIMENT_TREE_REMOVE,
-      ({ dvcRoot, id }: ExperimentItem) =>
-        this.experiments.runCommand(
-          AvailableCommands.EXPERIMENT_REMOVE,
-          dvcRoot,
-          id
-        )
+      async experimentItem => {
+        const selected = [...this.getSelectedExperimentItems(), experimentItem]
+
+        const deletable = collectDeletable(selected)
+
+        for (const [dvcRoot, ids] of Object.entries(deletable)) {
+          await this.experiments.runCommand(
+            AvailableCommands.EXPERIMENT_REMOVE,
+            dvcRoot,
+            ...ids
+          )
+        }
+      }
     )
   }
 
@@ -345,4 +338,8 @@ export class ExperimentsTree
   private getDisplayId(type: ExperimentType, label: string, id: string) {
     return type === ExperimentType.CHECKPOINT ? label : id
   }
+
+  private getSelectedExperimentItems() {
+    return [...this.view.selection]
+  }
 }

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

@@ -24,10 +24,7 @@ import plotsDiffFixture from '../../../fixtures/plotsDiff/output'
 import expShowFixture from '../../../fixtures/expShow/output'
 import { Operator } from '../../../../experiments/model/filterBy'
 import { joinColumnPath } from '../../../../experiments/columns/paths'
-import {
-  ExperimentItem,
-  ExperimentsTree
-} from '../../../../experiments/model/tree'
+import { ExperimentsTree } from '../../../../experiments/model/tree'
 import { buildExperiments, buildSingleRepoExperiments } from '../util'
 import { ResourceLocator } from '../../../../resourceLocator'
 import { InternalCommands } from '../../../../commands/internal'
@@ -42,6 +39,7 @@ import { Param } from '../../../../experiments/model/modify/collect'
 import { WorkspaceExperiments } from '../../../../experiments/workspace'
 import { ColumnType } from '../../../../experiments/webview/contract'
 import { copyOriginalColors } from '../../../../experiments/model/status/colors'
+import { ExperimentItem } from '../../../../experiments/model/collect'
 
 suite('Experiments Tree Test Suite', () => {
   const disposable = Disposable.fn()
@@ -465,16 +463,53 @@ suite('Experiments Tree Test Suite', () => {
     })
 
     it('should be able to remove an experiment with dvc.views.experimentsTree.removeExperiment', async () => {
+      const mockExperimentId = 'exp-to-remove'
+      const mockExperiment = {
+        dvcRoot: dvcDemoPath,
+        id: mockExperimentId,
+        type: ExperimentType.EXPERIMENT
+      }
+
+      const mockExperimentRemove = stub(
+        CliExecutor.prototype,
+        'experimentRemove'
+      ).resolves('')
+
+      stub(
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        (ExperimentsTree as any).prototype,
+        'getSelectedExperimentItems'
+      ).returns([mockExperiment])
+
+      await commands.executeCommand(
+        RegisteredCommands.EXPERIMENT_TREE_REMOVE,
+        mockExperiment
+      )
+
+      expect(mockExperimentRemove).to.be.calledWithExactly(
+        dvcDemoPath,
+        mockExperimentId
+      )
+    })
+
+    it('should be able to remove the provided experiment with dvc.views.experimentsTree.removeExperiment (if no experiments are selected)', async () => {
       const mockExperiment = 'exp-to-remove'
 
       const mockExperimentRemove = stub(
         CliExecutor.prototype,
         'experimentRemove'
       ).resolves('')
 
+      stub(
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        (ExperimentsTree as any).prototype,
+        'getSelectedExperimentItems'
+      ).returns([])
+
       await commands.executeCommand(RegisteredCommands.EXPERIMENT_TREE_REMOVE, {
         dvcRoot: dvcDemoPath,
-        id: mockExperiment
+        id: mockExperiment,
+        type: ExperimentType.EXPERIMENT
       })
 
       expect(mockExperimentRemove).to.be.calledWithExactly(
@@ -483,6 +518,51 @@ suite('Experiments Tree Test Suite', () => {
       )
     })
 
+    it('should be able to remove multiple experiments with dvc.views.experimentsTree.removeExperiment', async () => {
+      const mockExperimentId = 'exp-removed'
+      const mockQueuedExperimentLabel = 'queued-removed'
+
+      const mockExperimentRemove = stub(
+        CliExecutor.prototype,
+        'experimentRemove'
+      ).resolves('')
+
+      stub(
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        (ExperimentsTree as any).prototype,
+        'getSelectedExperimentItems'
+      ).returns([
+        dvcDemoPath,
+        {
+          dvcRoot: dvcDemoPath,
+          label: mockQueuedExperimentLabel,
+          type: ExperimentType.QUEUED
+        },
+        {
+          dvcRoot: dvcDemoPath,
+          id: 'checkpoint-excluded',
+          type: ExperimentType.CHECKPOINT
+        },
+        {
+          dvcRoot: dvcDemoPath,
+          id: 'workspace-excluded',
+          type: ExperimentType.WORKSPACE
+        }
+      ])
+
+      await commands.executeCommand(RegisteredCommands.EXPERIMENT_TREE_REMOVE, {
+        dvcRoot: dvcDemoPath,
+        id: mockExperimentId,
+        type: ExperimentType.EXPERIMENT
+      })
+
+      expect(mockExperimentRemove).to.be.calledWithExactly(
+        dvcDemoPath,
+        mockQueuedExperimentLabel,
+        mockExperimentId
+      )
+    })
+
     it('should be able to apply an experiment to the workspace with dvc.views.experimentsTree.applyExperiment', async () => {
       const mockExperiment = 'exp-to-apply'