data status: support targets to limit scope of command (#10789)

dvc/commands/data.py

@@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, ClassVar
 
 from funcy import chunks, compact, log_durations
 
-from dvc.cli import formatter
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -108,6 +108,7 @@ class CmdDataStatus(CmdBase):
     def run(self) -> int:
         with log_durations(logger.trace, "in data_status"):
             status = self.repo.data_status(
+                targets=self.args.targets,
                 granular=self.args.granular,
                 untracked_files=self.args.untracked_files,
                 not_in_remote=self.args.not_in_remote,
@@ -147,6 +148,14 @@ def add_parser(subparsers, parent_parser):
         formatter_class=formatter.RawDescriptionHelpFormatter,
         help=DATA_STATUS_HELP,
     )
+    data_status_parser.add_argument(
+        "targets",
+        nargs="*",
+        help=(
+            "Limit command scope to these tracked files/directories, "
+            ".dvc files and stage names."
+        ),
+    ).complete = completion.FILE  # type: ignore[attr-defined]
     data_status_parser.add_argument(
         "--json",
         action="store_true",

dvc/repo/data.py

@@ -1,17 +1,17 @@
 import os
 import posixpath
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any, TypedDict, Union
+from typing import TYPE_CHECKING, Optional, TypedDict, Union
 
 from dvc.fs.callbacks import DEFAULT_CALLBACK
-from dvc.repo.worktree import worktree_view
 from dvc.ui import ui
+from dvc_data.index.view import DataIndexView
 
 if TYPE_CHECKING:
     from dvc.fs.callbacks import Callback
     from dvc.repo import Repo
     from dvc.scm import Git, NoSCM
-    from dvc_data.index import DataIndex
+    from dvc_data.index import BaseDataIndex, DataIndex, DataIndexKey
     from dvc_data.index.diff import Change
 
 
@@ -47,12 +47,13 @@ def _adapt_path(change: "Change") -> str:
 
 
 def _diff(
-    old: "DataIndex",
-    new: "DataIndex",
+    old: "BaseDataIndex",
+    new: "BaseDataIndex",
     *,
     granular: bool = False,
     not_in_cache: bool = False,
     callback: "Callback" = DEFAULT_CALLBACK,
+    filter_keys: Optional[list["DataIndexKey"]] = None,
 ) -> dict[str, list[str]]:
     from dvc_data.index.diff import UNCHANGED, UNKNOWN, diff
 
@@ -74,6 +75,16 @@ def _diff(
         with_unknown=True,
         callback=callback,
     ):
+        # The index is a trie, so even when we filter by a specific path
+        # like `dir/file`, all parent nodes leading to that path (e.g., `dir/`)
+        # still appear in the view. As a result, keys like `dir/` will be present
+        # even if only `dir/file` matches the filter.
+        # We need to skip such entries to avoid showing root of tracked directories.
+        if filter_keys and not any(
+            change.key[: len(filter_key)] == filter_key for filter_key in filter_keys
+        ):
+            continue
+
         if (
             change.typ == UNCHANGED
             and (not change.old or not change.old.hash_info)
@@ -136,7 +147,45 @@ def _git_info(scm: Union["Git", "NoSCM"], untracked_files: str = "all") -> GitIn
     )
 
 
-def _diff_index_to_wtree(repo: "Repo", **kwargs: Any) -> dict[str, list[str]]:
+def filter_index(
+    index: Union["DataIndex", "DataIndexView"],
+    filter_keys: Optional[list["DataIndexKey"]] = None,
+) -> "BaseDataIndex":
+    if not filter_keys:
+        return index
+
+    if isinstance(index, DataIndexView):
+        orig_index = index._index
+        parent_filter_fn = index.filter_fn
+    else:
+        orig_index = index
+        parent_filter_fn = None
+
+    def filter_fn(key: "DataIndexKey") -> bool:
+        if parent_filter_fn is not None and not parent_filter_fn(key):
+            return False
+
+        for filter_key in filter_keys:
+            # eg: if key is "dir/file" and filter_key is "dir/", return True
+            if key[: len(filter_key)] == filter_key:
+                return True
+            # eg: if key is `dir/` and filter_key is `dir/file`, also return True.
+            # This ensures we include parent prefixes needed to reach matching leaves.
+            # Intermediate prefixes must be retained to access nested keys.
+            if filter_key[: len(key)] == key:
+                return True
+        return False
+
+    from dvc_data.index import view
+
+    return view(orig_index, filter_fn=filter_fn)
+
+
+def _diff_index_to_wtree(
+    repo: "Repo",
+    filter_keys: Optional[list["DataIndexKey"]] = None,
+    granular: bool = False,
+) -> dict[str, list[str]]:
     from .index import build_data_index
 
     with ui.progress(desc="Building workspace index", unit="entry") as pb:
@@ -147,30 +196,45 @@ def _diff_index_to_wtree(repo: "Repo", **kwargs: Any) -> dict[str, list[str]]:
             compute_hash=True,
             callback=pb.as_callback(),
         )
+        workspace_view = filter_index(workspace, filter_keys=filter_keys)
 
     with ui.progress(
         desc="Calculating diff between index/workspace",
         unit="entry",
     ) as pb:
+        index = repo.index.data["repo"]
+        view = filter_index(index, filter_keys=filter_keys)
         return _diff(
-            repo.index.data["repo"],
-            workspace,
+            view,
+            workspace_view,
+            filter_keys=filter_keys,
+            granular=granular,
             not_in_cache=True,
             callback=pb.as_callback(),
-            **kwargs,
         )
 
 
 def _diff_head_to_index(
-    repo: "Repo", head: str = "HEAD", **kwargs: Any
+    repo: "Repo",
+    head: str = "HEAD",
+    filter_keys: Optional[list["DataIndexKey"]] = None,
+    granular: bool = False,
 ) -> dict[str, list[str]]:
     index = repo.index.data["repo"]
+    index_view = filter_index(index, filter_keys=filter_keys)
 
     with repo.switch(head):
         head_index = repo.index.data["repo"]
+        head_view = filter_index(head_index, filter_keys=filter_keys)
 
     with ui.progress(desc="Calculating diff between head/index", unit="entry") as pb:
-        return _diff(head_index, index, callback=pb.as_callback(), **kwargs)
+        return _diff(
+            head_view,
+            index_view,
+            filter_keys=filter_keys,
+            granular=granular,
+            callback=pb.as_callback(),
+        )
 
 
 class Status(TypedDict):
@@ -204,46 +268,87 @@ def _transform_git_paths_to_dvc(repo: "Repo", files: Iterable[str]) -> list[str]
 
 def _get_entries_not_in_remote(
     repo: "Repo",
+    filter_keys: Optional[list["DataIndexKey"]] = None,
     granular: bool = False,
     remote_refresh: bool = False,
 ) -> list[str]:
     """Get entries that are not in remote storage."""
+    from dvc.repo.worktree import worktree_view
     from dvc_data.index import StorageKeyError
 
     # View into the index, with only pushable entries
-    view = worktree_view(repo.index, push=True).data["repo"]
+    index = worktree_view(repo.index, push=True)
+    data_index = index.data["repo"]
+
+    view = filter_index(data_index, filter_keys=filter_keys)  # type: ignore[arg-type]
 
     missing_entries = []
+    with ui.progress(desc="Checking remote", unit="entry") as pb:
         for key, entry in view.iteritems(shallow=not granular):
             if not (entry and entry.hash_info):
                 continue
 
+            # The index is a trie, so even when we filter by a specific path
+            # like `dir/file`, all parent nodes leading to that path (e.g., `dir/`)
+            # still appear in the view. As a result, keys like `dir/` will be present
+            # even if only `dir/file` matches the filter.
+            # We need to skip such entries to avoid showing root of tracked directories.
+            if filter_keys and not any(
+                key[: len(filter_key)] == filter_key for filter_key in filter_keys
+            ):
+                continue
+
             k = (*key, "") if entry.meta and entry.meta.isdir else key
             try:
                 if not view.storage_map.remote_exists(entry, refresh=remote_refresh):
                     missing_entries.append(os.path.sep.join(k))
+                    pb.update()
             except StorageKeyError:
                 pass
 
     return missing_entries
 
 
+def _matches_target(p: str, targets: Iterable[str]) -> bool:
+    sep = os.sep
+    return any(p == t or p.startswith(t + sep) for t in targets)
+
+
+def _prune_keys(filter_keys: list["DataIndexKey"]) -> list["DataIndexKey"]:
+    sorted_keys = sorted(set(filter_keys), key=len)
+    result: list[DataIndexKey] = []
+
+    for key in sorted_keys:
+        if not any(key[: len(prefix)] == prefix for prefix in result):
+            result.append(key)
+    return result
+
+
 def status(
     repo: "Repo",
+    targets: Optional[Iterable[Union[os.PathLike[str], str]]] = None,
+    *,
+    granular: bool = False,
     untracked_files: str = "no",
     not_in_remote: bool = False,
     remote_refresh: bool = False,
-    granular: bool = False,
     head: str = "HEAD",
 ) -> Status:
     from dvc.scm import NoSCMError, SCMError
 
-    uncommitted_diff = _diff_index_to_wtree(repo, granular=granular)
-    unchanged = set(uncommitted_diff.pop("unchanged", []))
+    targets = targets or []
+    filter_keys: list[DataIndexKey] = [repo.fs.relparts(os.fspath(t)) for t in targets]
+    # try to remove duplicate and overlapping keys
+    filter_keys = _prune_keys(filter_keys)
 
+    uncommitted_diff = _diff_index_to_wtree(
+        repo, filter_keys=filter_keys, granular=granular
+    )
+    unchanged = set(uncommitted_diff.pop("unchanged", []))
     entries_not_in_remote = (
         _get_entries_not_in_remote(
             repo,
+            filter_keys=filter_keys,
             granular=granular,
             remote_refresh=remote_refresh,
         )
@@ -252,16 +357,24 @@ def status(
     )
 
     try:
-        committed_diff = _diff_head_to_index(repo, head=head, granular=granular)
+        committed_diff = _diff_head_to_index(
+            repo, filter_keys=filter_keys, head=head, granular=granular
+        )
     except (SCMError, NoSCMError):
         committed_diff = {}
     else:
         unchanged &= set(committed_diff.pop("unchanged", []))
 
     git_info = _git_info(repo.scm, untracked_files=untracked_files)
-    untracked = git_info.get("untracked", [])
-    untracked = _transform_git_paths_to_dvc(repo, untracked)
-
+    scm_filter_targets = {
+        os.path.relpath(os.path.abspath(t), repo.scm.root_dir) for t in targets
+    }
+    untracked_it: Iterable[str] = git_info.get("untracked", [])
+    if scm_filter_targets:
+        untracked_it = (
+            f for f in untracked_it if _matches_target(f, scm_filter_targets)
+        )
+    untracked = _transform_git_paths_to_dvc(repo, untracked_it)
     # order matters here
     return Status(
         not_in_cache=uncommitted_diff.pop("not_in_cache", []),

tests/func/test_data_status.py

@@ -1,8 +1,11 @@
+from collections.abc import Iterable
+from functools import partial
 from os import fspath
 from os.path import join
 from typing import TYPE_CHECKING
 
 import pytest
+from pytest_test_utils import matchers as m
 
 from dvc.repo import Repo
 from dvc.repo.data import _transform_git_paths_to_dvc, posixpath_to_os_path
@@ -10,8 +13,6 @@ from dvc.testing.tmp_dir import TmpDir, make_subrepo
 from dvc.utils.fs import remove
 
 if TYPE_CHECKING:
-    from pytest_test_utils import matchers
-
     from dvc.stage import Stage
 
 EMPTY_STATUS = {
@@ -430,7 +431,7 @@ def test_missing_remote_cache(M, tmp_dir, dvc, scm, local_remote):
 
 
 def test_not_in_remote_respects_not_pushable(
-    M: type["matchers.Matcher"], tmp_dir: TmpDir, dvc: Repo, scm, mocker, local_remote
+    M: type["m.Matcher"], tmp_dir: TmpDir, dvc: Repo, scm, mocker, local_remote
 ):
     stages: list[Stage] = tmp_dir.dvc_gen({"foo": "foo", "dir": {"foobar": "foobar"}})
     # Make foo not pushable
@@ -535,3 +536,338 @@ def test_empty_dir(tmp_dir, scm, dvc, M):
         "uncommitted": {"modified": [join("data", "")]},
         "git": M.dict(),
     }
+
+
+def test_untracked_files_filter_targets(M, tmp_dir, scm, dvc):
+    tmp_dir.gen(
+        {"spam": "spam", "ham": "ham", "dir": {"eggs": "eggs", "bacon": "bacon"}}
+    )
+    _default = EMPTY_STATUS | {"git": M.dict()}
+    status = partial(dvc.data_status, untracked_files="all")
+
+    assert status(["not-existing"]) == _default
+
+    assert status(["spam"]) == _default | {"untracked": ["spam"]}
+    assert status(["spam", "ham"]) == _default | {
+        "untracked": M.unordered("spam", "ham")
+    }
+    assert status(["dir"]) == _default | {
+        "untracked": M.unordered(join("dir", "eggs"), join("dir", "bacon")),
+    }
+    assert status([join("dir", "")]) == _default | {
+        "untracked": M.unordered(join("dir", "eggs"), join("dir", "bacon")),
+    }
+    assert status([join("dir", "bacon")]) == _default | {
+        "untracked": [join("dir", "bacon")]
+    }
+
+
+def param(*values):
+    """Uses test id from the first value."""
+    first = values[0]
+    _id = (
+        ",".join(first)
+        if isinstance(first, Iterable) and not isinstance(first, str)
+        else first
+    )
+    return pytest.param(*values, id=_id)
+
+
+@pytest.mark.parametrize(
+    "targets,expected",
+    [
+        param(
+            ["foo"],
+            {"committed": {"added": ["foo"]}, "uncommitted": {"deleted": ["foo"]}},
+        ),
+        param(
+            ["bar"],
+            {"committed": {"added": ["bar"]}, "uncommitted": {"modified": ["bar"]}},
+        ),
+        param(["foobar"], {"committed": {"added": ["foobar"]}, "uncommitted": {}}),
+        param(["not-existing"], {}),
+        param(["baz"], {"untracked": ["baz"]}),
+        param(
+            ["foo", "foobar"],
+            {
+                "committed": {"added": m.unordered("foo", "foobar")},
+                "uncommitted": {"deleted": ["foo"]},
+            },
+        ),
+    ],
+)
+def test_filter_targets_files_after_dvc_commit(M, tmp_dir, dvc, scm, targets, expected):
+    tmp_dir.dvc_gen({"foo": "foo", "bar": "bar", "foobar": "foobar"})
+    (tmp_dir / "foo").unlink()  # deleted
+    tmp_dir.gen({"bar": "bar modified", "baz": "baz new"})
+
+    assert dvc.data_status(
+        targets=targets, untracked_files="all"
+    ) == EMPTY_STATUS | expected | {"git": M.dict()}
+    assert dvc.data_status(
+        targets=targets, granular=True, untracked_files="all"
+    ) == EMPTY_STATUS | expected | {"git": M.dict()}
+
+
+@pytest.mark.parametrize(
+    "targets,expected",
+    [
+        param(["not-existing"], {}),
+        param(["foo"], {"uncommitted": {"deleted": ["foo"]}}),
+        param(["bar"], {"unchanged": ["bar"]}),
+        param(["baz"], {"unchanged": ["baz"]}),
+        param(["foobar"], {"unchanged": ["foobar"]}),
+        param(
+            ("foo", "foobar"),
+            {"unchanged": ["foobar"], "uncommitted": {"deleted": ["foo"]}},
+        ),
+    ],
+)
+def test_filter_targets_after_git_commit(M, tmp_dir, dvc, scm, targets, expected):
+    tmp_dir.dvc_gen(
+        {"foo": "foo", "bar": "bar", "foobar": "foobar", "baz": "baz"},
+        commit="add files",
+    )
+    (tmp_dir / "foo").unlink()  # deleted
+
+    assert dvc.data_status(
+        targets=targets, untracked_files="all"
+    ) == EMPTY_STATUS | expected | {"git": M.dict()}
+    assert dvc.data_status(
+        targets=targets, granular=True, untracked_files="all"
+    ) == EMPTY_STATUS | expected | {"git": M.dict()}
+
+
+def with_aliases(values, aliases):
+    """Generate test cases by reusing values for given aliases from existing ones."""
+    for value in values:
+        targets = value[0]
+        assert isinstance(targets, tuple)
+        yield param(*value)
+    yield from (
+        param(alias, *rest)
+        for alias, to in aliases.items()
+        for targets, *rest in values
+        if to == targets
+    )
+
+
+@pytest.mark.parametrize(
+    "targets,expected_ng,expected_g",
+    with_aliases(
+        [
+            (
+                ("dir",),
+                {
+                    "committed": {"added": [join("dir", "")]},
+                    "uncommitted": {"modified": [join("dir", "")]},
+                },
+                {
+                    "committed": {
+                        "added": m.unordered(
+                            join("dir", ""),
+                            join("dir", "foo"),
+                            join("dir", "sub", "bar"),
+                            join("dir", "foobar"),
+                        )
+                    },
+                    "uncommitted": {
+                        "added": [join("dir", "baz")],
+                        "modified": [join("dir", ""), join("dir", "sub", "bar")],
+                        "deleted": [join("dir", "foo")],
+                    },
+                },
+            ),
+            (
+                (join("dir", "foo"),),
+                {},
+                {
+                    "committed": {"added": [join("dir", "foo")]},
+                    "uncommitted": {"deleted": [join("dir", "foo")]},
+                },
+            ),
+            (
+                (join("dir", "baz"),),
+                {},
+                {
+                    "uncommitted": {"added": [join("dir", "baz")]},
+                },
+            ),
+            (
+                (join("dir", "sub"),),
+                {},
+                {
+                    "committed": {"added": [join("dir", "sub", "bar")]},
+                    "uncommitted": {"modified": [join("dir", "sub", "bar")]},
+                },
+            ),
+            (
+                (join("dir", "sub", "bar"),),
+                {},
+                {
+                    "committed": {"added": [join("dir", "sub", "bar")]},
+                    "uncommitted": {"modified": [join("dir", "sub", "bar")]},
+                },
+            ),
+            (
+                (join("dir", "foobar"),),
+                {},
+                {
+                    "committed": {"added": [join("dir", "foobar")]},
+                    "uncommitted": {},
+                },
+            ),
+            ((join("dir", "not-existing-file"),), {}, {}),
+            ((join("dir", "not-existing-dir", ""),), {}, {}),
+            ((join("dir", "sub", "not-existing-file"),), {}, {}),
+            (
+                (join("dir", "foo"), join("dir", "foobar")),
+                {},
+                {
+                    "committed": {
+                        "added": m.unordered(join("dir", "foo"), join("dir", "foobar"))
+                    },
+                    "uncommitted": {"deleted": [join("dir", "foo")]},
+                },
+            ),
+        ],
+        {
+            # the values for these are used from above test cases
+            (join("dir", ""),): ("dir",),
+            (join("dir", "sub", ""),): (join("dir", "sub"),),
+            (join("dir", ""), join("dir", "foo")): ("dir",),
+        },
+    ),
+)
+def test_filter_targets_inside_directory_after_dvc_commit(
+    M, tmp_dir, dvc, scm, targets, expected_ng, expected_g
+):
+    tmp_dir.dvc_gen({"dir": {"foo": "foo", "sub": {"bar": "bar"}, "foobar": "foobar"}})
+    (tmp_dir / "dir" / "foo").unlink()  # deleted
+    (tmp_dir / "dir" / "sub" / "bar").write_text("bar modified")
+    (tmp_dir / "dir" / "baz").write_text("baz new")
+
+    assert dvc.data_status(
+        targets=targets, untracked_files="all"
+    ) == EMPTY_STATUS | expected_ng | {"git": M.dict()}
+    assert dvc.data_status(
+        targets=targets, granular=True, untracked_files="all"
+    ) == EMPTY_STATUS | expected_g | {"git": M.dict()}
+
+
+@pytest.mark.parametrize(
+    "targets,expected_ng,expected_g",
+    with_aliases(
+        [
+            (
+                (join("dir", "foo"),),
+                {},
+                {"committed": {"deleted": [join("dir", "foo")]}},
+            ),
+            ((join("dir", "baz"),), {}, {"committed": {"added": [join("dir", "baz")]}}),
+            ((join("dir", "foobar"),), {}, {"unchanged": [join("dir", "foobar")]}),
+            (
+                ("dir",),
+                {"committed": {"modified": [join("dir", "")]}},
+                {
+                    "unchanged": [join("dir", "foobar")],
+                    "committed": {
+                        "added": [join("dir", "baz")],
+                        "modified": m.unordered(
+                            join("dir", ""), join("dir", "sub", "bar")
+                        ),
+                        "deleted": [join("dir", "foo")],
+                    },
+                },
+            ),
+            (
+                (join("dir", "sub"),),
+                {},
+                {"committed": {"modified": [join("dir", "sub", "bar")]}},
+            ),
+            (
+                (join("dir", "sub", "bar"),),
+                {},
+                {"committed": {"modified": [join("dir", "sub", "bar")]}},
+            ),
+            (
+                (join("dir", "foo"), join("dir", "foobar")),
+                {},
+                {
+                    "unchanged": [join("dir", "foobar")],
+                    "committed": {"deleted": [join("dir", "foo")]},
+                },
+            ),
+        ],
+        {
+            (join("dir", ""),): ("dir",),
+            (join("dir", "sub", ""),): (join("dir", "sub"),),
+            (join("dir", ""), join("dir", "foo")): ("dir",),
+        },
+    ),
+)
+def test_filter_targets_inside_directory_after_git_commit(
+    M, tmp_dir, dvc, scm, targets, expected_ng, expected_g
+):
+    tmp_dir.dvc_gen(
+        {"dir": {"foo": "foo", "sub": {"bar": "bar"}, "foobar": "foobar"}},
+        commit="add dir",
+    )
+    (tmp_dir / "dir" / "foo").unlink()  # deleted
+    (tmp_dir / "dir" / "sub" / "bar").write_text("bar modified")
+    (tmp_dir / "dir" / "baz").write_text("baz new")
+    dvc.add(["dir"])
+
+    assert dvc.data_status(
+        targets=targets, untracked_files="all"
+    ) == EMPTY_STATUS | expected_ng | {"git": M.dict()}
+    assert dvc.data_status(
+        targets=targets, granular=True, untracked_files="all"
+    ) == EMPTY_STATUS | expected_g | {"git": M.dict()}
+
+
+@pytest.mark.parametrize("to_check", ["remote", "cache"])
+@pytest.mark.parametrize(
+    "targets, non_granular, granular",
+    [
+        param(("foo",), ["foo"], ["foo"]),
+        param(("dir",), [join("dir", "")], [join("dir", ""), join("dir", "bar")]),
+        param(
+            (join("dir", ""),), [join("dir", "")], [join("dir", ""), join("dir", "bar")]
+        ),
+        param((join("dir", "bar"),), [], [join("dir", "bar")]),
+        param(
+            (join("dir", "bar"), "foo"),
+            ["foo"],
+            m.unordered(join("dir", "bar"), "foo"),
+        ),
+        param(
+            (join("dir", "bar"), "dir"),
+            [join("dir", "")],
+            m.unordered(join("dir", ""), join("dir", "bar")),
+        ),
+        param(
+            ("dir", "foo"),
+            m.unordered(join("dir", ""), "foo"),
+            m.unordered(join("dir", ""), join("dir", "bar"), "foo"),
+        ),
+    ],
+)
+def test_filter_targets_not_in_cache(
+    M, local_remote, tmp_dir, scm, dvc, to_check, targets, non_granular, granular
+):
+    tmp_dir.dvc_gen({"foo": "foo", "dir": {"bar": "bar"}})
+
+    if to_check == "cache":
+        dvc.push()
+        dvc.cache.local.clear()
+
+    not_in_remote = to_check == "remote"
+    key = "not_in_" + to_check
+    d = EMPTY_STATUS | {"git": M.dict(), "committed": M.dict()}
+    assert dvc.data_status(targets, not_in_remote=not_in_remote) == d | {
+        key: non_granular
+    }
+    assert dvc.data_status(targets, granular=True, not_in_remote=not_in_remote) == d | {
+        key: granular
+    }

tests/unit/command/test_data_status.py

@@ -50,6 +50,7 @@ def test_cli(dvc, mocker, mocked_status):
     cmd = cli_args.func(cli_args)
     assert cmd.run() == 0
     status.assert_called_once_with(
+        targets=[],
         untracked_files="all",
         not_in_remote=False,
         remote_refresh=True,