Skip to content

Commit

Permalink
worktree: use progress bars in push/fetch
Browse files Browse the repository at this point in the history
  • Loading branch information
pmrowla committed Nov 23, 2022
1 parent 2748492 commit af255a4
Showing 1 changed file with 20 additions and 2 deletions.
22 changes: 20 additions & 2 deletions dvc/repo/worktree.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from typing import TYPE_CHECKING, Optional

from dvc.fs.callbacks import Callback

if TYPE_CHECKING:
from dvc.cloud import Remote
from dvc.index import Index, IndexView
Expand Down Expand Up @@ -57,7 +59,12 @@ def fetch_worktree(repo: "Repo", remote: "Remote") -> int:
remote.path,
*key,
)
return save(index)
total = len(index)
with Callback.as_tqdm_callback(
unit="file", desc="Fetch", disable=total == 0
) as cb:
cb.set_size(total)
return save(index, callback=cb)


def push_worktree(repo: "Repo", remote: "Remote") -> int:
Expand All @@ -66,7 +73,18 @@ def push_worktree(repo: "Repo", remote: "Remote") -> int:

view = worktree_view(repo.index, push=True)
index = view.data["repo"]
pushed = checkout(index, remote.path, remote.fs, latest_only=False)
total = len(index)
with Callback.as_tqdm_callback(
unit="file", desc="Push", disable=total == 0
) as cb:
cb.set_size(total)
pushed = checkout(
index,
remote.path,
remote.fs,
latest_only=False,
callback=cb,
)

for stage in view.stages:
for out in stage.outs:
Expand Down

0 comments on commit af255a4

Please sign in to comment.