diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index bd869e9b7..4a44d4f66 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -85,12 +85,14 @@ class TempDirExecutor(BaseLocalExecutor): def _init_git(self, scm: "Git", branch: Optional[str] = None, **kwargs): from dulwich.repo import Repo as DulwichRepo + from ..utils import push_refspec + DulwichRepo.init(os.fspath(self.root_dir)) refspec = f"{EXEC_NAMESPACE}/" - scm.push_refspec(self.git_url, refspec, refspec, **kwargs) + push_refspec(scm, self.git_url, refspec, refspec, **kwargs) if branch: - scm.push_refspec(self.git_url, branch, branch, **kwargs) + push_refspec(scm, self.git_url, branch, branch, **kwargs) self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True) elif self.scm.get_ref(EXEC_BRANCH): self.scm.remove_ref(EXEC_BRANCH) diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py index 5cb480504..d7c085460 100644 --- a/dvc/repo/experiments/executor/ssh.py +++ b/dvc/repo/experiments/executor/ssh.py @@ -104,6 +104,8 @@ class SSHExecutor(BaseExecutor): return kwargs def _init_git(self, scm: "Git", branch: Optional[str] = None, **kwargs): + from ..utils import push_refspec + with self.sshfs() as fs: fs.makedirs(self.root_dir) self._ssh_cmd(fs, "git init .") @@ -120,9 +122,9 @@ class SSHExecutor(BaseExecutor): # (see https://github.com/iterative/dvc/issues/6508) kwargs.update(self._git_client_args(fs)) refspec = f"{EXEC_NAMESPACE}/" - scm.push_refspec(self.git_url, refspec, refspec, **kwargs) + push_refspec(scm, self.git_url, refspec, refspec, **kwargs) if branch: - scm.push_refspec(self.git_url, branch, branch, **kwargs) + push_refspec(scm, self.git_url, branch, branch, **kwargs) self._ssh_cmd(fs, f"git symbolic-ref {EXEC_BRANCH} {branch}") else: self._ssh_cmd( diff --git a/dvc/repo/experiments/pull.py b/dvc/repo/experiments/pull.py index 1b7f157e0..cd649d820 100644 --- a/dvc/repo/experiments/pull.py +++ b/dvc/repo/experiments/pull.py @@ -31,9 +31,17 @@ def pull( refspec = f"{exp_ref}:{exp_ref}" logger.debug("git pull experiment '%s' -> '%s'", git_remote, refspec) - repo.scm.fetch_refspecs( - git_remote, [refspec], force=force, on_diverged=on_diverged - ) + + from dvc.scm import TqdmGit + + with TqdmGit(desc="Fetching git refs") as pbar: + repo.scm.fetch_refspecs( + git_remote, + [refspec], + force=force, + on_diverged=on_diverged, + progress=pbar.update_git, + ) if pull_cache: _pull_cache(repo, exp_ref, **kwargs) diff --git a/dvc/repo/experiments/push.py b/dvc/repo/experiments/push.py index eec925395..6c7764444 100644 --- a/dvc/repo/experiments/push.py +++ b/dvc/repo/experiments/push.py @@ -37,14 +37,19 @@ def push( refname = str(exp_ref) logger.debug("git push experiment '%s' -> '%s'", exp_ref, git_remote) - push_refspec( - repo.scm, - git_remote, - refname, - refname, - force=force, - on_diverged=on_diverged, - ) + + from dvc.scm import TqdmGit + + with TqdmGit(desc="Pushing git refs") as pbar: + push_refspec( + repo.scm, + git_remote, + refname, + refname, + force=force, + on_diverged=on_diverged, + progress=pbar.update_git, + ) if push_cache: _push_cache(repo, exp_ref, **kwargs) diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index 7508a338f..5bb41955c 100644 --- a/dvc/repo/experiments/remove.py +++ b/dvc/repo/experiments/remove.py @@ -80,8 +80,17 @@ def _remove_commited_exps( if not remote: remove_exp_refs(repo.scm, remove_list) else: + from dvc.scm import TqdmGit + for ref_info in remove_list: - push_refspec(repo.scm, remote, None, str(ref_info)) + with TqdmGit(desc="Pushing git refs") as pbar: + push_refspec( + repo.scm, + remote, + None, + str(ref_info), + progress=pbar.update_git, + ) return remain_list diff --git a/dvc/scm/__init__.py b/dvc/scm/__init__.py index f1fe39e10..ef26017f4 100644 --- a/dvc/scm/__init__.py +++ b/dvc/scm/__init__.py @@ -1,10 +1,14 @@ """Manages source control systems (e.g. Git).""" from contextlib import contextmanager -from typing import Iterator +from typing import TYPE_CHECKING, Iterator +from dvc.progress import Tqdm from dvc.scm.base import Base, NoSCMError from dvc.scm.git import Git +if TYPE_CHECKING: + from dvc.scm.progress import GitProgressEvent + # Syntactic sugar to signal that this is an actual implementation for a DVC # project under no SCM control. @@ -50,14 +54,30 @@ def SCM( ) +class TqdmGit(Tqdm): + def __init__(self, *args, **kwargs): + kwargs.setdefault("unit", "obj") + super().__init__(*args, **kwargs) + + def update_git(self, event: "GitProgressEvent") -> None: + phase, completed, total, message, *_ = event + if phase: + message = (phase + " | " + message) if message else phase + if message: + self.postfix["info"] = f" {message} | " + if completed: + self.update_to(completed, total) + + def clone(url: str, to_path: str, **kwargs): from .base import CloneError from .exceptions import CloneError as InternalCloneError - try: - return Git.clone(url, to_path, **kwargs) - except InternalCloneError as exc: - raise CloneError(str(exc)) + with TqdmGit(desc="Cloning") as pbar: + try: + return Git.clone(url, to_path, progress=pbar.update_git, **kwargs) + except InternalCloneError as exc: + raise CloneError(str(exc)) def resolve_rev(scm: "Git", rev: str) -> str: diff --git a/dvc/scm/git/backend/base.py b/dvc/scm/git/backend/base.py index 9f9ae91bc..b01e95d1c 100644 --- a/dvc/scm/git/backend/base.py +++ b/dvc/scm/git/backend/base.py @@ -15,6 +15,8 @@ from dvc.scm.exceptions import SCMError from ..objects import GitObject if TYPE_CHECKING: + from dvc.scm.progress import GitProgressEvent + from ..objects import GitCommit @@ -49,6 +51,7 @@ class BaseGitBackend(ABC): to_path: str, rev: Optional[str] = None, shallow_branch: Optional[str] = None, + progress: Callable[["GitProgressEvent"], None] = None, ): pass @@ -197,6 +200,7 @@ class BaseGitBackend(ABC): dest: str, force: bool = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + progress: Callable[["GitProgressEvent"], None] = None, **kwargs, ): """Push refspec to a remote Git repo. @@ -223,6 +227,7 @@ class BaseGitBackend(ABC): refspecs: Iterable[str], force: Optional[bool] = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + progress: Callable[["GitProgressEvent"], None] = None, **kwargs, ): """Fetch refspecs from a remote Git repo. diff --git a/dvc/scm/git/backend/dulwich/__init__.py b/dvc/scm/git/backend/dulwich/__init__.py index 6a7176160..c41cb42a2 100644 --- a/dvc/scm/git/backend/dulwich/__init__.py +++ b/dvc/scm/git/backend/dulwich/__init__.py @@ -18,7 +18,6 @@ from typing import ( from funcy import cached_property -from dvc.progress import Tqdm from dvc.scm.exceptions import AuthError, InvalidRemote, SCMError from dvc.utils import relpath @@ -26,10 +25,12 @@ from ...objects import GitObject from ..base import BaseGitBackend if TYPE_CHECKING: + from dvc.scm.progress import GitProgressEvent from dvc.types import StrPath from ...objects import GitCommit + logger = logging.getLogger(__name__) @@ -135,6 +136,7 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method to_path: str, rev: Optional[str] = None, shallow_branch: Optional[str] = None, + progress: Callable[["GitProgressEvent"], None] = None, ): raise NotImplementedError @@ -396,6 +398,7 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method dest: str, force: bool = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + progress: Callable[["GitProgressEvent"], None] = None, **kwargs, ): from dulwich.client import HTTPUnauthorized, get_transport_and_path @@ -439,22 +442,14 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method return new_refs try: - with Tqdm( - desc="Pushing git refs", bar_format=self.BAR_FMT_NOTOTAL - ) as pbar: + from dvc.scm.progress import GitProgressReporter - def progress(msg_b): - msg = msg_b.decode("ascii").strip() - pbar.update_msg(msg) - pbar.refresh() - logger.trace(msg) - - client.send_pack( - path, - update_refs, - self.repo.object_store.generate_pack_data, - progress=progress, - ) + client.send_pack( + path, + update_refs, + self.repo.object_store.generate_pack_data, + progress=GitProgressReporter(progress) if progress else None, + ) except (NotGitRepository, SendPackError) as exc: raise SCMError("Git failed to push '{src}' to '{url}'") from exc except HTTPUnauthorized: @@ -484,6 +479,7 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method refspecs: Iterable[str], force: Optional[bool] = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + progress: Callable[["GitProgressEvent"], None] = None, **kwargs, ): from dulwich.client import get_transport_and_path @@ -519,22 +515,14 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method f"'{url}' is not a valid Git remote or URL" ) from exc - with Tqdm( - desc="Fetching git refs", bar_format=self.BAR_FMT_NOTOTAL - ) as pbar: + from dvc.scm.progress import GitProgressReporter - def progress(msg_b): - msg = msg_b.decode("ascii").strip() - pbar.update_msg(msg) - pbar.refresh() - logger.trace(msg) - - fetch_result = client.fetch( - path, - self.repo, - progress=progress, - determine_wants=determine_wants, - ) + fetch_result = client.fetch( + path, + self.repo, + progress=GitProgressReporter(progress) if progress else None, + determine_wants=determine_wants, + ) for (lh, rh, _) in fetch_refs: try: if rh in self.repo.refs: diff --git a/dvc/scm/git/backend/gitpython.py b/dvc/scm/git/backend/gitpython.py index 430538758..0a6d1e3da 100644 --- a/dvc/scm/git/backend/gitpython.py +++ b/dvc/scm/git/backend/gitpython.py @@ -16,7 +16,6 @@ from typing import ( from funcy import ignore -from dvc.progress import Tqdm from dvc.scm.exceptions import ( CloneError, MergeConflictError, @@ -29,43 +28,12 @@ from ..objects import GitCommit, GitObject from .base import BaseGitBackend if TYPE_CHECKING: + from dvc.scm.progress import GitProgressEvent from dvc.types import StrPath logger = logging.getLogger(__name__) -class TqdmGit(Tqdm): - def update_git(self, op_code, cur_count, max_count=None, message=""): - op_code = self.code2desc(op_code) - if op_code: - message = (op_code + " | " + message) if message else op_code - if message: - self.postfix["info"] = f" {message} | " - self.update_to(cur_count, max_count) - - @staticmethod - def code2desc(op_code): - from git import RootUpdateProgress as OP - - ops = { - OP.COUNTING: "Counting", - OP.COMPRESSING: "Compressing", - OP.WRITING: "Writing", - OP.RECEIVING: "Receiving", - OP.RESOLVING: "Resolving", - OP.FINDING_SOURCES: "Finding sources", - OP.CHECKING_OUT: "Checking out", - OP.CLONE: "Cloning", - OP.FETCH: "Fetching", - OP.UPDWKTREE: "Updating working tree", - OP.REMOVE: "Removing", - OP.PATHCHANGE: "Changing path", - OP.URLCHANGE: "Changing URL", - OP.BRANCHCHANGE: "Changing branch", - } - return ops.get(op_code & OP.OP_MASK, "") - - class GitPythonObject(GitObject): def __init__(self, obj): self.obj = obj @@ -152,6 +120,7 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method to_path: str, rev: Optional[str] = None, shallow_branch: Optional[str] = None, + progress: Callable[["GitProgressEvent"], None] = None, ): from git import Repo from git.exc import GitCommandError @@ -174,19 +143,23 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method # git disables --depth for local clones unless file:// url # scheme is used url = f"file://{url}" - with TqdmGit(desc="Cloning", unit="obj") as pbar: - clone_from = partial( - Repo.clone_from, - url, - to_path, - env=env, # needed before we can fix it in __init__ - no_single_branch=True, - progress=pbar.update_git, - ) - if shallow_branch is None: - tmp_repo = clone_from() - else: - tmp_repo = clone_from(branch=shallow_branch, depth=1) + + from dvc.scm.progress import GitProgressReporter + + clone_from = partial( + Repo.clone_from, + url, + to_path, + env=env, # needed before we can fix it in __init__ + no_single_branch=True, + progress=GitProgressReporter.wrap_fn(progress) + if progress + else None, + ) + if shallow_branch is None: + tmp_repo = clone_from() + else: + tmp_repo = clone_from(branch=shallow_branch, depth=1) tmp_repo.close() except GitCommandError as exc: # pylint: disable=no-member raise CloneError(url, to_path) from exc @@ -471,6 +444,7 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method dest: str, force: bool = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + progress: Callable[["GitProgressEvent"], None] = None, **kwargs, ): raise NotImplementedError @@ -481,6 +455,7 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method refspecs: Iterable[str], force: Optional[bool] = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + progress: Callable[["GitProgressEvent"], None] = None, **kwargs, ): raise NotImplementedError diff --git a/dvc/scm/git/backend/pygit2.py b/dvc/scm/git/backend/pygit2.py index a973da176..5ccc04189 100644 --- a/dvc/scm/git/backend/pygit2.py +++ b/dvc/scm/git/backend/pygit2.py @@ -27,9 +27,9 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: + from dvc.scm.progress import GitProgressEvent from dvc.types import StrPath - # NOTE: constant from libgit2 git2/checkout.h # This can be removed after next pygit2 release: # see https://github.com/libgit2/pygit2/pull/1087 @@ -153,6 +153,7 @@ class Pygit2Backend(BaseGitBackend): # pylint:disable=abstract-method to_path: str, rev: Optional[str] = None, shallow_branch: Optional[str] = None, + progress: Callable[["GitProgressEvent"], None] = None, ): raise NotImplementedError @@ -374,6 +375,7 @@ class Pygit2Backend(BaseGitBackend): # pylint:disable=abstract-method dest: str, force: bool = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + progress: Callable[["GitProgressEvent"], None] = None, **kwargs, ): raise NotImplementedError @@ -384,6 +386,7 @@ class Pygit2Backend(BaseGitBackend): # pylint:disable=abstract-method refspecs: Iterable[str], force: Optional[bool] = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + progress: Callable[["GitProgressEvent"], None] = None, **kwargs, ): raise NotImplementedError diff --git a/dvc/scm/progress.py b/dvc/scm/progress.py new file mode 100644 index 000000000..e4d306c6e --- /dev/null +++ b/dvc/scm/progress.py @@ -0,0 +1,73 @@ +from typing import NamedTuple, Optional, Union + +from funcy import compose + + +def code2desc(op_code): + from git import RootUpdateProgress as OP + + ops = { + OP.COUNTING: "Counting", + OP.COMPRESSING: "Compressing", + OP.WRITING: "Writing", + OP.RECEIVING: "Receiving", + OP.RESOLVING: "Resolving", + OP.FINDING_SOURCES: "Finding sources", + OP.CHECKING_OUT: "Checking out", + OP.CLONE: "Cloning", + OP.FETCH: "Fetching", + OP.UPDWKTREE: "Updating working tree", + OP.REMOVE: "Removing", + OP.PATHCHANGE: "Changing path", + OP.URLCHANGE: "Changing URL", + OP.BRANCHCHANGE: "Changing branch", + } + return ops.get(op_code & OP.OP_MASK, "") + + +class GitProgressEvent(NamedTuple): + phase: str = "" + completed: Optional[int] = None + total: Optional[int] = None + message: str = "" + + @classmethod + def parsed_from_gitpython( + cls, + op_code, + cur_count, + max_count=None, + message="", # pylint: disable=redefined-outer-name + ): + return cls(code2desc(op_code), cur_count, max_count, message) + + +class GitProgressReporter: + def __init__(self, fn) -> None: + from git.util import CallableRemoteProgress + + self._reporter = CallableRemoteProgress(self.wrap_fn(fn)) + + def __call__(self, msg: Union[str, bytes]) -> None: + self._reporter._parse_progress_line( + msg.decode("ascii").strip() if isinstance(msg, bytes) else msg + ) + + @staticmethod + def wrap_fn(fn): + return compose(fn, GitProgressEvent.parsed_from_gitpython) + + +if __name__ == "__main__": + message = """ +Cloning into 'dvcyaml-schema'... +remote: Enumerating objects: 76, done. +remote: Counting objects: 100% (76/76), done. +remote: Compressing objects: 100% (56/56), done. +remote: Total 76 (delta 30), reused 38 (delta 9), pack-reused 0 +Receiving objects: 100% (76/76), 24.83 KiB | 89.00 KiB/s, done. +Resolving deltas: 100% (30/30), done.""" + + reporter = GitProgressReporter(print) + for line in message.splitlines(): + reporter(line) diff --git a/tests/func/test_external_repo.py b/tests/func/test_external_repo.py index 2305f23cb..9b65f0463 100644 --- a/tests/func/test_external_repo.py +++ b/tests/func/test_external_repo.py @@ -140,7 +140,9 @@ def test_shallow_clone_branch(erepo_dir): with repo.open_by_relpath("file") as fd: assert fd.read() == "branch" - mock_clone.assert_called_with(url, ANY, shallow_branch="branch") + mock_clone.assert_called_with( + url, ANY, shallow_branch="branch", progress=ANY + ) _, shallow = CLONES[url] assert shallow @@ -166,7 +168,9 @@ def test_shallow_clone_tag(erepo_dir): with repo.open_by_relpath("file") as fd: assert fd.read() == "foo" - mock_clone.assert_called_with(url, ANY, shallow_branch="v1") + mock_clone.assert_called_with( + url, ANY, shallow_branch="v1", progress=ANY + ) _, shallow = CLONES[url] assert shallow