feat: ignore files not in remote when push is false (#10749)

Co-authored-by: Saugat Pachhai (सौगात) <[email protected]>

dvc/repo/data.py

@@ -4,6 +4,7 @@ from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any, TypedDict, Union
 
 from dvc.fs.callbacks import DEFAULT_CALLBACK
+from dvc.repo.worktree import worktree_view
 from dvc.ui import ui
 
 if TYPE_CHECKING:
@@ -51,11 +52,8 @@ def _diff(
     *,
     granular: bool = False,
     not_in_cache: bool = False,
-    not_in_remote: bool = False,
-    remote_refresh: bool = False,
     callback: "Callback" = DEFAULT_CALLBACK,
 ) -> dict[str, list[str]]:
-    from dvc_data.index import StorageError
     from dvc_data.index.diff import UNCHANGED, UNKNOWN, diff
 
     ret: dict[str, list[str]] = {}
@@ -97,19 +95,6 @@ def _diff(
             # NOTE: emulating previous behaviour
             _add_change("not_in_cache", change)
 
-        try:
-            if (
-                not_in_remote
-                and change.old
-                and change.old.hash_info
-                and not old.storage_map.remote_exists(
-                    change.old, refresh=remote_refresh
-                )
-            ):
-                _add_change("not_in_remote", change)
-        except StorageError:
-            pass
-
         _add_change(change.typ, change)
 
     return ret
@@ -217,15 +202,57 @@ def _transform_git_paths_to_dvc(repo: "Repo", files: Iterable[str]) -> list[str]
     return [repo.fs.relpath(file, start) for file in files]
 
 
-def status(repo: "Repo", untracked_files: str = "no", **kwargs: Any) -> Status:
+def _get_entries_not_in_remote(
+    repo: "Repo",
+    granular: bool = False,
+    remote_refresh: bool = False,
+) -> list[str]:
+    """Get entries that are not in remote storage."""
+    from dvc_data.index import StorageKeyError
+
+    # View into the index, with only pushable entries
+    view = worktree_view(repo.index, push=True).data["repo"]
+
+    missing_entries = []
+    for key, entry in view.iteritems(shallow=not granular):
+        if not (entry and entry.hash_info):
+            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))
+        except StorageKeyError:
+            pass
+
+    return missing_entries
+
+
+def status(
+    repo: "Repo",
+    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
 
-    head = kwargs.pop("head", "HEAD")
-    uncommitted_diff = _diff_index_to_wtree(repo, **kwargs)
+    uncommitted_diff = _diff_index_to_wtree(repo, granular=granular)
     unchanged = set(uncommitted_diff.pop("unchanged", []))
 
+    entries_not_in_remote = (
+        _get_entries_not_in_remote(
+            repo,
+            granular=granular,
+            remote_refresh=remote_refresh,
+        )
+        if not_in_remote
+        else []
+    )
+
     try:
-        committed_diff = _diff_head_to_index(repo, head=head, **kwargs)
+        committed_diff = _diff_head_to_index(repo, head=head, granular=granular)
     except (SCMError, NoSCMError):
         committed_diff = {}
     else:
@@ -234,10 +261,11 @@ def status(repo: "Repo", untracked_files: str = "no", **kwargs: Any) -> Status:
     git_info = _git_info(repo.scm, untracked_files=untracked_files)
     untracked = git_info.get("untracked", [])
     untracked = _transform_git_paths_to_dvc(repo, untracked)
+
     # order matters here
     return Status(
         not_in_cache=uncommitted_diff.pop("not_in_cache", []),
-        not_in_remote=uncommitted_diff.pop("not_in_remote", []),
+        not_in_remote=entries_not_in_remote,
         committed=committed_diff,
         uncommitted=uncommitted_diff,
         untracked=untracked,

tests/func/test_data_status.py

@@ -1,13 +1,19 @@
 from os import fspath
 from os.path import join
+from typing import TYPE_CHECKING
 
 import pytest
 
 from dvc.repo import Repo
 from dvc.repo.data import _transform_git_paths_to_dvc, posixpath_to_os_path
-from dvc.testing.tmp_dir import make_subrepo
+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 = {
     "committed": {},
     "uncommitted": {},
@@ -423,6 +429,51 @@ 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
+):
+    stages: list[Stage] = tmp_dir.dvc_gen({"foo": "foo", "dir": {"foobar": "foobar"}})
+    # Make foo not pushable
+    stages[0].outs[0].can_push = False
+    stages[0].dump()
+
+    def assert_not_in_remote_is(
+        granular: bool, not_in_remote: list[str], committed: list[str]
+    ):
+        assert dvc.data_status(
+            granular=granular, remote_refresh=True, not_in_remote=True
+        ) == {
+            **EMPTY_STATUS,
+            "git": M.dict(),
+            "not_in_remote": M.unordered(*not_in_remote),
+            "committed": {"added": M.unordered(*committed)},
+        }
+
+    foo = "foo"
+    dir_ = join("dir", "")
+    foobar = join("dir", "foobar")
+
+    assert_not_in_remote_is(
+        granular=True,
+        not_in_remote=[dir_, foobar],
+        committed=[foo, dir_, foobar],
+    )
+    assert_not_in_remote_is(granular=False, not_in_remote=[dir_], committed=[foo, dir_])
+
+    dvc.push()
+
+    assert_not_in_remote_is(
+        granular=True,
+        not_in_remote=[],
+        committed=[foo, dir_, foobar],
+    )
+    assert_not_in_remote_is(
+        granular=False,
+        not_in_remote=[],
+        committed=[foo, dir_],
+    )
+
+
 def test_root_from_dir_to_file(M, tmp_dir, dvc, scm):
     tmp_dir.dvc_gen({"data": {"foo": "foo", "bar": "bar"}})
     remove("data")