Add multi-select push and pull to tracked explorer tree (#1809)

* add multi select push and pull actions

* add unit test for new functionality

* extend unit tests

* add test for siblings being returned

* standardise tree view name

* have path item properties present in the source control view

* set state of repository dependants after experiments data is processed

* mock correct file

Author: mattseddon

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]
+    })
+  })
+})

extension/src/repository/model/collect.ts

@@ -135,12 +135,16 @@ export const collectTrackedOuts = (data: ExperimentsOutput): Set<string> => {
 }
 
 export const collectTrackedPaths = async (
-  { dvcRoot, resourceUri, isTracked }: PathItem,
+  pathItem: string | PathItem,
   getChildren: (path: string) => Promise<PathItem[]>
 ): Promise<string[]> => {
-  const acc = []
+  const acc: string[] = []
+  if (typeof pathItem === 'string') {
+    return acc
+  }
 
-  if (isTracked) {
+  const { dvcRoot, resourceUri, isTracked } = pathItem
+  if (isTracked !== false) {
     acc.push(relativeWithUri(dvcRoot, resourceUri))
     return acc
   }
@@ -151,6 +155,89 @@ export const collectTrackedPaths = async (
   return acc
 }
 
+type SelectedPathAccumulator = { [dvcRoot: string]: (string | PathItem)[] }
+
+const collectSelectedPaths = (pathItems: (string | PathItem)[]): string[] => {
+  const acc = new Set<string>()
+  for (const pathItem of pathItems) {
+    if (typeof pathItem === 'string') {
+      acc.add(pathItem)
+      continue
+    }
+    acc.add(pathItem.resourceUri.fsPath)
+  }
+  return [...acc]
+}
+
+const parentIsSelected = (fsPath: string, paths: string[]) => {
+  for (const path of paths.filter(path => path !== fsPath)) {
+    if (fsPath.includes(path)) {
+      return true
+    }
+  }
+  return false
+}
+
+const initializeAccumulatorRoot = (
+  acc: SelectedPathAccumulator,
+  dvcRoot: string
+) => {
+  if (!acc[dvcRoot]) {
+    acc[dvcRoot] = []
+  }
+}
+
+const collectPathItem = (
+  acc: SelectedPathAccumulator,
+  addedPaths: Set<string>,
+  pathItem: PathItem | string,
+  dvcRoot: string,
+  path: string
+) => {
+  initializeAccumulatorRoot(acc, dvcRoot)
+  acc[dvcRoot].push(pathItem)
+  addedPaths.add(path)
+}
+
+const collectRoot = (
+  acc: SelectedPathAccumulator,
+  addedPaths: Set<string>,
+  path: string
+) => collectPathItem(acc, addedPaths, path, path, path)
+
+const collectRootOrPathItem = (
+  acc: SelectedPathAccumulator,
+  addedPaths: Set<string>,
+  paths: string[],
+  pathItem: string | PathItem
+) => {
+  const path =
+    typeof pathItem === 'string' ? pathItem : pathItem.resourceUri.fsPath
+  if (addedPaths.has(path) || parentIsSelected(path, paths)) {
+    return
+  }
+
+  if (typeof pathItem === 'string') {
+    collectRoot(acc, addedPaths, path)
+    return
+  }
+
+  const { dvcRoot } = pathItem
+  collectPathItem(acc, addedPaths, pathItem, dvcRoot, path)
+}
+
+export const collectSelected = (
+  pathItems: (string | PathItem)[]
+): SelectedPathAccumulator => {
+  const selectedPaths = collectSelectedPaths(pathItems)
+  const acc: SelectedPathAccumulator = {}
+  const addedPaths = new Set<string>()
+  for (const pathItem of pathItems) {
+    collectRootOrPathItem(acc, addedPaths, selectedPaths, pathItem)
+  }
+  return acc
+}
+
 const isUncollectedChild = (
   deleted: Set<string>,
   deletedPath: string,