From cf976f17aed569beeca287aaeb99d990a2cb83d6 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Mon, 26 Jun 2023 00:10:44 +0300
Subject: [PATCH] data: status: use ui.progress instead of ui.status (#9662)

With all fetch and checkout related things that we've done in the recent past,
we now have all the mechanisms necessary to provide actual progress (no total,
but still better than spinner) during `dvc data status` operations.
---
 dvc/repo/data.py                       | 28 +++++++++++++++++++++-----
 tests/unit/command/test_data_status.py |  6 +++---
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/dvc/repo/data.py b/dvc/repo/data.py
index 209d3b293..b36548675 100644
--- a/dvc/repo/data.py
+++ b/dvc/repo/data.py
@@ -2,9 +2,11 @@ import os
 import posixpath
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, TypedDict, Union
 
+from dvc.fs.callbacks import DEFAULT_CALLBACK
 from dvc.ui import ui
 
 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
@@ -50,6 +52,7 @@ def _diff(
     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
@@ -70,6 +73,7 @@ def _diff(
         shallow=not granular,
         hash_only=True,
         with_unknown=True,
+        callback=callback,
     ):
         if (
             change.typ == UNCHANGED
@@ -149,16 +153,27 @@ 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]]:
     from .index import build_data_index
 
-    with ui.status("Building workspace index"):
+    with ui.progress(
+        desc="Building workspace index",
+        unit="entry",
+    ) as pb:
         workspace = build_data_index(
-            repo.index, repo.root_dir, repo.fs, compute_hash=True
+            repo.index,
+            repo.root_dir,
+            repo.fs,
+            compute_hash=True,
+            callback=pb.as_callback(),
         )
 
-    with ui.status("Calculating diff between index/workspace"):
+    with ui.progress(
+        desc="Calculating diff between index/workspace",
+        unit="entry",
+    ) as pb:
         return _diff(
             repo.index.data["repo"],
             workspace,
             not_in_cache=True,
+            callback=pb.as_callback(),
             **kwargs,
         )
 
@@ -171,8 +186,11 @@ def _diff_head_to_index(
     with repo.switch(head):
         head_index = repo.index.data["repo"]
 
-    with ui.status("Calculating diff between head/index"):
-        return _diff(head_index, index, **kwargs)
+    with ui.progress(
+        desc="Calculating diff between head/index",
+        unit="entry",
+    ) as pb:
+        return _diff(head_index, index, callback=pb.as_callback(), **kwargs)
 
 
 class Status(TypedDict):
diff --git a/tests/unit/command/test_data_status.py b/tests/unit/command/test_data_status.py
index 1e7e8cb3e..f37d270a6 100644
--- a/tests/unit/command/test_data_status.py
+++ b/tests/unit/command/test_data_status.py
@@ -75,7 +75,8 @@ def test_json(dvc, mocker, capsys, mocked_status, args, to_omit):
 
 def test_no_changes_repo(dvc, scm, capsys):
     assert main(["data", "status"]) == 0
-    assert capsys.readouterr() == ("No changes.\n", "")
+    out, _ = capsys.readouterr()
+    assert out == "No changes.\n"
 
 
 def test_empty_scm_repo(tmp_dir, capsys):
@@ -83,7 +84,7 @@ def test_empty_scm_repo(tmp_dir, capsys):
     Repo.init()
 
     assert main(["data", "status"]) == 0
-    out, err = capsys.readouterr()
+    out, _ = capsys.readouterr()
     assert (
         out
         == """\
@@ -91,7 +92,6 @@ No changes in an empty git repo.
 (there are changes not tracked by dvc, use "git status" to see)
 """
     )
-    assert not err
 
 
 @pytest.mark.parametrize(