Merge branch 'main' into sticky-exp-column

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/fileSystem/tree.ts

@@ -4,13 +4,13 @@ import {
   TreeDataProvider,
   TreeItem,
   TreeItemCollapsibleState,
-  Uri,
-  window
+  TreeView,
+  Uri
 } from 'vscode'
 import { exists, relativeWithUri } from '.'
 import { fireWatcher } from './watcher'
 import { deleteTarget, moveTargets } from './workspace'
-import { definedAndNonEmpty } from '../util/array'
+import { definedAndNonEmpty, uniqueValues } from '../util/array'
 import {
   AvailableCommands,
   CommandId,
@@ -26,16 +26,22 @@ import { warnOfConsequences } from '../vscode/modal'
 import { Response } from '../vscode/response'
 import { Resource } from '../repository/commands'
 import { WorkspaceRepositories } from '../repository/workspace'
-import { collectTrackedPaths, PathItem } from '../repository/model/collect'
+import {
+  collectSelected,
+  collectTrackedPaths,
+  PathItem
+} from '../repository/model/collect'
 import { Title } from '../vscode/title'
 import { Disposable } from '../class/dispose'
+import { createTreeView } from '../vscode/tree'
 
 export class TrackedExplorerTree
   extends Disposable
   implements TreeDataProvider<PathItem>
 {
   public readonly onDidChangeTreeData: Event<void>
 
+  private readonly view: TreeView<string | PathItem>
   private readonly internalCommands: InternalCommands
   private readonly repositories: WorkspaceRepositories
 
@@ -57,8 +63,8 @@ export class TrackedExplorerTree
 
     this.onDidChangeTreeData = repositories.treeDataChanged.event
 
-    this.dispose.track(
-      window.registerTreeDataProvider('dvc.views.trackedExplorerTree', this)
+    this.view = this.dispose.track(
+      createTreeView<PathItem>('dvc.views.trackedExplorerTree', this, true)
     )
   }
 
@@ -242,13 +248,29 @@ export class TrackedExplorerTree
 
   private tryThenForce(commandId: CommandId) {
     return async (pathItem: PathItem) => {
-      const { dvcRoot } = pathItem
-      const tracked = await collectTrackedPaths(pathItem, (path: string) =>
-        this.getRepoChildren(dvcRoot, path)
-      )
-      const args = [dvcRoot, ...tracked.sort()]
+      const selected = collectSelected([
+        ...this.getSelectedPathItems(),
+        pathItem
+      ])
+
+      for (const [dvcRoot, pathItems] of Object.entries(selected)) {
+        const tracked = []
+        for (const pathItem of pathItems) {
+          tracked.push(
+            ...(await collectTrackedPaths(pathItem, (path: string) =>
+              this.getRepoChildren(dvcRoot, path)
+            ))
+          )
+        }
 
-      return tryThenMaybeForce(this.internalCommands, commandId, ...args)
+        const args = [dvcRoot, ...uniqueValues(tracked).sort()]
+
+        await tryThenMaybeForce(this.internalCommands, commandId, ...args)
+      }
     }
   }
+
+  private getSelectedPathItems() {
+    return [...this.view.selection]
+  }
 }

extension/src/repository/index.ts

@@ -69,6 +69,7 @@ export class Repository extends DeferredDisposable {
       experiments.onDidChangeExperiments(data => {
         if (data) {
           this.model.transformAndSetExperiments(data)
+          this.setState()
         }
       })
     )

extension/src/repository/model/collect.test.ts

@@ -1,11 +1,12 @@
 import { join } from 'path'
 import { Uri } from 'vscode'
-import { collectTree } from './collect'
+import { collectSelected, collectTree } from './collect'
 import { dvcDemoPath } from '../../test/util'
 
+const makeUri = (...paths: string[]): Uri =>
+  Uri.file(join(dvcDemoPath, ...paths))
+
 describe('collectTree', () => {
-  const makeUri = (...paths: string[]): Uri =>
-    Uri.file(join(dvcDemoPath, ...paths))
   const makeAbsPath = (...paths: string[]): string => makeUri(...paths).fsPath
 
   it('should transform recursive list output into a tree', () => {
@@ -190,3 +191,87 @@ describe('collectTree', () => {
     )
   })
 })
+
+describe('collectSelected', () => {
+  const logsPathItem = {
+    dvcRoot: dvcDemoPath,
+    isDirectory: true,
+    isTracked: true,
+    resourceUri: makeUri('logs')
+  }
+
+  const accPathItem = {
+    dvcRoot: dvcDemoPath,
+    isDirectory: false,
+    isTracked: true,
+    resourceUri: makeUri('logs', 'acc.tsv')
+  }
+
+  const lossPathItem = {
+    dvcRoot: dvcDemoPath,
+    isDirectory: false,
+    isTracked: true,
+    resourceUri: makeUri('logs', 'loss.tsv')
+  }
+
+  it('should return an empty object if no path items are provided', () => {
+    expect(collectSelected([])).toStrictEqual({})
+  })
+
+  it('should return the original item if only one is provided', () => {
+    const selected = collectSelected([logsPathItem])
+
+    expect(selected).toStrictEqual({
+      [dvcDemoPath]: [logsPathItem]
+    })
+  })
+
+  it('should return a root given it is select', () => {
+    const selected = collectSelected([dvcDemoPath, logsPathItem, accPathItem])
+
+    expect(selected).toStrictEqual({
+      [dvcDemoPath]: [dvcDemoPath]
+    })
+  })
+
+  it('should return siblings if a parent is not provided', () => {
+    const selected = collectSelected([accPathItem, lossPathItem])
+
+    expect(selected).toStrictEqual({
+      [dvcDemoPath]: [accPathItem, lossPathItem]
+    })
+  })
+
+  it('should exclude all children from the final list', () => {
+    const selected = collectSelected([lossPathItem, accPathItem, logsPathItem])
+
+    expect(selected).toStrictEqual({
+      [dvcDemoPath]: [logsPathItem]
+    })
+  })
+
+  it('should return multiple entries when multiple roots are provided', () => {
+    const mockOtherRepoItem = {
+      dvcRoot: __dirname,
+      isDirectory: true,
+      isTracked: true,
+      resourceUri: Uri.file(join(__dirname, 'mock', 'path'))
+    }
+
+    const selected = collectSelected([
+      mockOtherRepoItem,
+      {
+        dvcRoot: dvcDemoPath,
+        isDirectory: false,
+        isTracked: true,
+        resourceUri: makeUri('logs', 'acc.tsv')
+      },
+      logsPathItem
+    ])
+
+    expect(selected).toStrictEqual({
+      [__dirname]: [mockOtherRepoItem],
+      [dvcDemoPath]: [logsPathItem]
+    })
+  })
+})