fix checkout callback during add/commit (iterative#9445)

* fix checkout callback

Fix iterative#6226.

* bump minversion of dvc-data to 0.49.1

dvc/output.py

@@ -36,7 +36,7 @@
 
 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 @@
     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 @@ def commit(self, filter_info=None, relink=True) -> None:
                     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 @@ def checkout(
         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 @@ def checkout(
                 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 @@ def add(  # noqa: C901
         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
 

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(

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",