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

@@ -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,6 +739,10 @@ class Output:
                     hardlink=hardlink,
                 )
             if relink:
+                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,
@@ -748,6 +751,7 @@ class Output:
                         relink=True,
                         state=self.repo.state,
                         prompt=prompt.confirm,
+                        progress_callback=callback,
                     )
                 self.set_exec()
 
@@ -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,6 +1356,10 @@ class Output:
         otransfer(staging, self.cache, {obj.hash_info}, hardlink=relink, shallow=False)
 
         if relink:
+            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,
@@ -1350,6 +1368,7 @@ class Output:
                     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",