From c7f797407a12fcaba657f2caaefd32f4de7265b0 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 1 Jul 2022 11:00:53 +1000
Subject: [PATCH] Only pull selected files if they include the invoked path

---
 extension/src/fileSystem/tree.ts              |  5 +-
 .../src/repository/model/collect.test.ts      | 61 ++++++++++++++++---
 extension/src/repository/model/collect.ts     | 20 +++++-
 3 files changed, 70 insertions(+), 16 deletions(-)

diff --git a/extension/src/fileSystem/tree.ts b/extension/src/fileSystem/tree.ts
index ba1d8c5fdc..941bedda65 100644
--- a/extension/src/fileSystem/tree.ts
+++ b/extension/src/fileSystem/tree.ts
@@ -258,10 +258,7 @@ export class TrackedExplorerTree
 
   private tryThenForce(commandId: CommandId) {
     return async (pathItem: PathItem) => {
-      const selected = collectSelected([
-        ...this.getSelectedPathItems(),
-        pathItem
-      ])
+      const selected = collectSelected(pathItem, this.getSelectedPathItems())
 
       for (const [dvcRoot, pathItems] of Object.entries(selected)) {
         const tracked = []
diff --git a/extension/src/repository/model/collect.test.ts b/extension/src/repository/model/collect.test.ts
index f2eb1587b1..b7a0c7aba7 100644
--- a/extension/src/repository/model/collect.test.ts
+++ b/extension/src/repository/model/collect.test.ts
@@ -214,20 +214,28 @@ describe('collectSelected', () => {
     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 no other items are selected', () => {
+    const selected = collectSelected(logsPathItem, [])
+
+    expect(selected).toStrictEqual({
+      [dvcDemoPath]: [logsPathItem]
+    })
   })
 
-  it('should return the original item if only one is provided', () => {
-    const selected = collectSelected([logsPathItem])
+  it('should return only the invoked item if it is not included in the selected paths', () => {
+    const selected = collectSelected(lossPathItem, [accPathItem])
 
     expect(selected).toStrictEqual({
-      [dvcDemoPath]: [logsPathItem]
+      [dvcDemoPath]: [lossPathItem]
     })
   })
 
   it('should return a root given it is select', () => {
-    const selected = collectSelected([dvcDemoPath, logsPathItem, accPathItem])
+    const selected = collectSelected(logsPathItem, [
+      dvcDemoPath,
+      logsPathItem,
+      accPathItem
+    ])
 
     expect(selected).toStrictEqual({
       [dvcDemoPath]: [dvcDemoPath]
@@ -235,7 +243,7 @@ describe('collectSelected', () => {
   })
 
   it('should return siblings if a parent is not provided', () => {
-    const selected = collectSelected([accPathItem, lossPathItem])
+    const selected = collectSelected(lossPathItem, [accPathItem, lossPathItem])
 
     expect(selected).toStrictEqual({
       [dvcDemoPath]: [accPathItem, lossPathItem]
@@ -243,13 +251,25 @@ describe('collectSelected', () => {
   })
 
   it('should exclude all children from the final list', () => {
-    const selected = collectSelected([lossPathItem, accPathItem, logsPathItem])
+    const selected = collectSelected(logsPathItem, [
+      lossPathItem,
+      accPathItem,
+      logsPathItem
+    ])
 
     expect(selected).toStrictEqual({
       [dvcDemoPath]: [logsPathItem]
     })
   })
 
+  it('should not exclude children from the final list if the invoked item is not in the selected list', () => {
+    const selected = collectSelected(accPathItem, [lossPathItem, logsPathItem])
+
+    expect(selected).toStrictEqual({
+      [dvcDemoPath]: [accPathItem]
+    })
+  })
+
   it('should return multiple entries when multiple roots are provided', () => {
     const mockOtherRepoItem = {
       dvcRoot: __dirname,
@@ -258,7 +278,7 @@ describe('collectSelected', () => {
       resourceUri: Uri.file(join(__dirname, 'mock', 'path'))
     }
 
-    const selected = collectSelected([
+    const selected = collectSelected(logsPathItem, [
       mockOtherRepoItem,
       {
         dvcRoot: dvcDemoPath,
@@ -274,4 +294,27 @@ describe('collectSelected', () => {
       [dvcDemoPath]: [logsPathItem]
     })
   })
+
+  it('should only return the invoked item if multiple roots are provided but the invoked item is not selected', () => {
+    const mockOtherRepoItem = {
+      dvcRoot: __dirname,
+      isDirectory: true,
+      isTracked: true,
+      resourceUri: Uri.file(join(__dirname, 'mock', 'path'))
+    }
+
+    const selected = collectSelected(logsPathItem, [
+      mockOtherRepoItem,
+      {
+        dvcRoot: dvcDemoPath,
+        isDirectory: false,
+        isTracked: true,
+        resourceUri: makeUri('logs', 'acc.tsv')
+      }
+    ])
+
+    expect(selected).toStrictEqual({
+      [dvcDemoPath]: [logsPathItem]
+    })
+  })
 })
diff --git a/extension/src/repository/model/collect.ts b/extension/src/repository/model/collect.ts
index fb07c1b901..0a2f545971 100644
--- a/extension/src/repository/model/collect.ts
+++ b/extension/src/repository/model/collect.ts
@@ -227,12 +227,26 @@ const collectRootOrPathItem = (
 }
 
 export const collectSelected = (
-  pathItems: (string | PathItem)[]
+  invokedPathItem: PathItem,
+  selectedPathItems: (string | PathItem)[]
 ): SelectedPathAccumulator => {
-  const selectedPaths = collectSelectedPaths(pathItems)
+  const invokedPath = invokedPathItem.resourceUri.fsPath
+
+  if (
+    !selectedPathItems.some(pathItem => {
+      if (typeof pathItem === 'string') {
+        return pathItem === invokedPath
+      }
+      return pathItem.resourceUri.fsPath === invokedPath
+    })
+  ) {
+    return { [invokedPathItem.dvcRoot]: [invokedPathItem] }
+  }
+
+  const selectedPaths = collectSelectedPaths(selectedPathItems)
   const acc: SelectedPathAccumulator = {}
   const addedPaths = new Set<string>()
-  for (const pathItem of pathItems) {
+  for (const pathItem of selectedPathItems) {
     collectRootOrPathItem(acc, addedPaths, selectedPaths, pathItem)
   }
   return acc