From 628b99244283ba92cfd173abc29d71f979f7a4bb Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Tue, 8 Nov 2022 17:56:55 +0900
Subject: [PATCH] index: fix granular dir output view filtering

---
 dvc/repo/index.py             | 51 ++++++++++++++++++++---------------
 tests/func/test_repo_index.py | 23 ++++++++++++++++
 2 files changed, 52 insertions(+), 22 deletions(-)

diff --git a/dvc/repo/index.py b/dvc/repo/index.py
index 94fbe635f..a4273465f 100644
--- a/dvc/repo/index.py
+++ b/dvc/repo/index.py
@@ -7,8 +7,10 @@ from typing import (
     Iterable,
     Iterator,
     List,
+    NamedTuple,
     Optional,
     Set,
+    Tuple,
 )
 
 from funcy import cached_property, nullcontext
@@ -348,6 +350,11 @@ class Index:
         return dict_md5(self.dumpd())
 
 
+class _DataPrefixes(NamedTuple):
+    explicit: Set["DataIndexKey"]
+    recursive: Set["DataIndexKey"]
+
+
 class IndexView:
     """Read-only view of Index.data using filtered stages."""
 
@@ -391,33 +398,32 @@ class IndexView:
             yield from stage.deps
 
     @property
-    def outs(self) -> Iterator["Output"]:
-        outs = set()
+    def _filtered_outs(self) -> Iterator[Tuple["Output", Optional[str]]]:
         for stage, filter_info in self._stage_infos:
             for out in stage.filter_outs(filter_info):
                 if not self._outs_filter or self._outs_filter(out):
-                    outs.add(out)
-        yield from outs
+                    yield out, filter_info
+
+    @property
+    def outs(self) -> Iterator["Output"]:
+        yield from {out for (out, _) in self._filtered_outs}
 
     @cached_property
-    def _data_prefixes(self) -> Dict[str, Set["DataIndexKey"]]:
+    def _data_prefixes(self) -> Dict[str, "_DataPrefixes"]:
         from collections import defaultdict
 
-        prefixes: Dict[str, Set["DataIndexKey"]] = defaultdict(set)
-        for stage, filter_info in self._stage_infos:
-            for out in stage.filter_outs(filter_info):
-                if self._outs_filter and not self._outs_filter(out):
-                    continue
-                workspace, key = out.index_key
-                if filter_info and out.fs.path.isin(filter_info, out.fs_path):
-                    key = (
-                        *key,
-                        out.fs.path.relparts(filter_info, out.fs_path),
-                    )
-                prefixes[workspace].add(key)
-                prefixes[workspace].update(
-                    key[:i] for i in range(len(key), 0, -1)
-                )
+        prefixes: Dict[str, "_DataPrefixes"] = defaultdict(
+            lambda: _DataPrefixes(set(), set())
+        )
+        for out, filter_info in self._filtered_outs:
+            workspace, key = out.index_key
+            if filter_info and out.fs.path.isin(filter_info, out.fs_path):
+                key = key + out.fs.path.relparts(filter_info, out.fs_path)
+            if out.meta.isdir:
+                prefixes[workspace].recursive.add(key)
+            prefixes[workspace].explicit.update(
+                key[:i] for i in range(len(key), 0, -1)
+            )
         return prefixes
 
     @cached_property
@@ -428,9 +434,10 @@ class IndexView:
 
         def key_filter(workspace: str, key: "DataIndexKey"):
             try:
-                return key in self._data_prefixes[workspace] or any(
+                prefixes = self._data_prefixes[workspace]
+                return key in prefixes.explicit or any(
                     key[: len(prefix)] == prefix
-                    for prefix in self._data_prefixes[workspace]
+                    for prefix in prefixes.recursive
                 )
             except KeyError:
                 return False
diff --git a/tests/func/test_repo_index.py b/tests/func/test_repo_index.py
index 0056c97aa..c3dbec448 100644
--- a/tests/func/test_repo_index.py
+++ b/tests/func/test_repo_index.py
@@ -305,6 +305,29 @@ def test_getitem(tmp_dir, dvc, run_copy):
         _ = index["no-valid-stage-name"]
 
 
+def test_view_granular_dir(tmp_dir, scm, dvc, run_copy):
+    tmp_dir.dvc_gen(
+        {"dir": {"subdir": {"in-subdir": "in-subdir"}, "in-dir": "in-dir"}},
+        commit="init",
+    )
+    index = Index(dvc)
+
+    # view should include the specific target, parent dirs, and children
+    # view should exclude any siblings of the target
+    view = index.targets_view("dir/subdir")
+    data_index = view.data["repo"]
+    assert ("dir",) in data_index
+    assert (
+        "dir",
+        "subdir",
+    ) in data_index
+    assert ("dir", "subdir", "in-subdir") in data_index
+    assert (
+        "dir",
+        "in-dir",
+    ) not in data_index
+
+
 def test_view_stage_filter(tmp_dir, scm, dvc, run_copy):
     (stage1,) = tmp_dir.dvc_gen("foo", "foo")
     stage2 = run_copy("foo", "bar", name="copy-foo-bar")