From d08c85f7d11c8c9c3433299c6821858c546f1c42 Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Fri, 12 May 2023 18:14:13 +0545
Subject: [PATCH] fix checkout callback during add/commit (#9445)

* fix checkout callback

Fix https://github.com/iterative/dvc/issues/6226.

* bump minversion of dvc-data to 0.49.1
---
 dvc/output.py        | 65 ++++++++++++++++++++++++++++----------------
 dvc/repo/worktree.py |  3 +-
 pyproject.toml       |  2 +-
 3 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/dvc/output.py b/dvc/output.py
index 6a70b23a9..e6295e233 100644
--- a/dvc/output.py
+++ b/dvc/output.py
@@ -36,7 +36,7 @@ from dvc_objects.errors import ObjectFormatError
 
 from .annotations import ANNOTATION_FIELDS, ANNOTATION_SCHEMA, Annotation
 from .fs import LocalFileSystem, RemoteMissingDepsError, Schemes, get_cloud_fs
-from .fs.callbacks import DEFAULT_CALLBACK
+from .fs.callbacks import DEFAULT_CALLBACK, Callback
 from .utils import relpath
 from .utils.fs import path_isin
 
@@ -45,7 +45,6 @@ if TYPE_CHECKING:
     from dvc_data.index import DataIndexKey
     from dvc_objects.db import ObjectDB
 
-    from .fs.callbacks import Callback
     from .ignore import DvcIgnoreFilter
 
 logger = logging.getLogger(__name__)
@@ -740,15 +739,20 @@ class Output:
                     hardlink=hardlink,
                 )
             if relink:
-                self._checkout(
-                    filter_info or self.fs_path,
-                    self.fs,
-                    obj,
-                    self.cache,
-                    relink=True,
-                    state=self.repo.state,
-                    prompt=prompt.confirm,
-                )
+                rel = self.fs.path.relpath(filter_info or self.fs_path)
+                with Callback.as_tqdm_callback(
+                    desc=f"Checking out {rel}", unit="files"
+                ) as callback:
+                    self._checkout(
+                        filter_info or self.fs_path,
+                        self.fs,
+                        obj,
+                        self.cache,
+                        relink=True,
+                        state=self.repo.state,
+                        prompt=prompt.confirm,
+                        progress_callback=callback,
+                    )
                 self.set_exec()
 
     def _commit_granular_dir(self, filter_info, hardlink) -> Optional["HashFile"]:
@@ -890,9 +894,19 @@ class Output:
         checkpoint_reset: bool = False,
         **kwargs,
     ) -> Optional[Tuple[bool, Optional[bool]]]:
+        # callback passed act as a aggregate callback.
+        # do not let checkout to call set_size and change progressbar.
+        class CallbackProxy(Callback):
+            def relative_update(self, inc: int = 1) -> None:
+                progress_callback.relative_update(inc)
+                return super().relative_update(inc)
+
+            def branch(self, *args, **kwargs):
+                return progress_callback.branch(*args, **kwargs)
+
+        callback = CallbackProxy()
         if not self.use_cache:
-            if progress_callback != DEFAULT_CALLBACK:
-                progress_callback.relative_update(self.get_files_number(filter_info))
+            callback.relative_update(self.get_files_number(filter_info))
             return None
 
         obj = self.get_obj(filter_info=filter_info)
@@ -914,7 +928,7 @@ class Output:
                 obj,
                 self.cache,
                 force=force,
-                progress_callback=progress_callback,
+                progress_callback=callback,
                 relink=relink,
                 state=self.repo.state,
                 prompt=prompt.confirm,
@@ -1342,15 +1356,20 @@ class Output:
         otransfer(staging, self.cache, {obj.hash_info}, hardlink=relink, shallow=False)
 
         if relink:
-            self._checkout(
-                path,
-                self.fs,
-                obj,
-                self.cache,
-                relink=True,
-                state=self.repo.state,
-                prompt=prompt.confirm,
-            )
+            rel = self.fs.path.relpath(path)
+            with Callback.as_tqdm_callback(
+                desc=f"Checking out {rel}", unit="files"
+            ) as callback:
+                self._checkout(
+                    path,
+                    self.fs,
+                    obj,
+                    self.cache,
+                    relink=True,
+                    state=self.repo.state,
+                    prompt=prompt.confirm,
+                    progress_callback=callback,
+                )
             self.set_exec()
         return obj
 
diff --git a/dvc/repo/worktree.py b/dvc/repo/worktree.py
index 6caf14cb1..45fc6dff3 100644
--- a/dvc/repo/worktree.py
+++ b/dvc/repo/worktree.py
@@ -173,7 +173,7 @@ def push_worktree(
         ) as cb:
             cb.set_size(total)
             try:
-                pushed += checkout(
+                stats = checkout(
                     new_index,
                     remote_obj.path,
                     remote_obj.fs,
@@ -184,6 +184,7 @@ def push_worktree(
                     jobs=jobs,
                     **diff_kwargs,
                 )
+                pushed += sum(len(changes) for changes in stats.values())
             except VersioningNotSupported:
                 logger.exception("")
                 raise DvcException(
diff --git a/pyproject.toml b/pyproject.toml
index 231a9b19b..ec8ca2be5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,7 +34,7 @@ dependencies = [
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=0.47.5,<0.48",
+    "dvc-data>=0.49.1,<0.50",
     "dvc-http>=2.29.0",
     "dvc-render>=0.3.1,<1",
     "dvc-studio-client>=0.9.0,<1",