fix(data-status): iterate/expand index upto tracked directory in shallow mode (#10900)

dvc/repo/data.py

@@ -1,6 +1,6 @@
 import os
 import posixpath
-from collections import defaultdict
+from collections import defaultdict, deque
 from collections.abc import Iterable, Iterator, Mapping
 from typing import TYPE_CHECKING, Optional, TypedDict, Union
 
@@ -385,6 +385,26 @@ def _transform_git_paths_to_dvc(repo: "Repo", files: Iterable[str]) -> list[str]
     return [repo.fs.relpath(file, start) for file in files]
 
 
+def iter_index(
+    index: Union["BaseDataIndex", "DataIndexView"], shallow: bool = False
+) -> Iterator[tuple["DataIndexKey", Optional["DataIndexEntry"]]]:
+    if not shallow:
+        yield from index.iteritems(shallow=shallow)
+        return
+
+    # only iterate until we find entries with hash_info in shallow mode
+    todo: deque[tuple[DataIndexKey, dict]] = deque([((), index.info(()))])
+    while todo:
+        key, info = todo.popleft()
+        entry = info.get("entry")
+        if info.get("type") == "directory" and not (entry and entry.hash_info):
+            try:
+                todo.extend(index.ls(key, detail=True))
+            except (KeyError, DataIndexDirError):
+                pass
+        yield key, entry
+
+
 def _get_entries_not_in_remote(
     repo: "Repo",
     filter_keys: Optional[Iterable["DataIndexKey"]] = None,
@@ -418,7 +438,7 @@ def _get_entries_not_in_remote(
 
     n = 0
     with TqdmCallback(size=n, desc="Checking remote", unit="entry") as cb:
-        for key, entry in view.iteritems(shallow=not granular):
+        for key, entry in iter_index(view, shallow=not granular):
             if not (entry and entry.hash_info):
                 continue
 

tests/func/test_data_status.py

@@ -1211,3 +1211,40 @@ def test_renames(
         dvc.data_status(granular=True, with_renames=True, targets=targets)
         == EMPTY_STATUS | {"git": M.dict()} | expected_granular
     )
+
+
+def test_shallow_should_iterate_upto_tracked_directory(
+    M, tmp_dir, dvc, scm, local_remote
+):
+    """Testing regression for https://github.com/iterative/dvc/issues/10899."""
+
+    tmp_dir.scm_gen({"dir": {".gitkeep": ""}}, commit="mk dir")
+    dataset_root = join("dir", "data", "")
+    tmp_dir.dvc_gen({dataset_root: {"foo": "foo", "bar": "bar"}})
+    files = [dataset_root, join(dataset_root, "foo"), join(dataset_root, "bar")]
+
+    assert dvc.data_status(not_in_remote=True) == {
+        **EMPTY_STATUS,
+        "committed": {"added": [dataset_root]},
+        "not_in_remote": [dataset_root],
+        "git": M.dict(),
+    }
+    assert dvc.data_status(granular=True, not_in_remote=True) == {
+        **EMPTY_STATUS,
+        "committed": {"added": M.unordered(*files)},
+        "not_in_remote": M.unordered(*files),
+        "git": M.dict(),
+    }
+
+    assert dvc.push() == 3
+
+    assert dvc.data_status(not_in_remote=True, remote_refresh=True) == {
+        **EMPTY_STATUS,
+        "committed": {"added": [dataset_root]},
+        "git": M.dict(),
+    }
+    assert dvc.data_status(granular=True, not_in_remote=True, remote_refresh=True) == {
+        **EMPTY_STATUS,
+        "committed": {"added": M.unordered(*files)},
+        "git": M.dict(),
+    }