From 351b952d5fdc153b2f97808442dfeacdfbe9cf47 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Wed, 8 Mar 2023 16:54:30 +0900 Subject: [PATCH] brancher: add brancher.switch helper --- dvc/repo/__init__.py | 16 ++++++- dvc/repo/brancher.py | 84 ++++++++++++++++++++++++++++------- dvc/repo/experiments/utils.py | 7 ++- dvc/scm.py | 3 +- 4 files changed, 90 insertions(+), 20 deletions(-) diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 1b6525ad3..a0aaadd67 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -3,7 +3,16 @@ import os from collections import defaultdict from contextlib import contextmanager from functools import wraps -from typing import TYPE_CHECKING, Callable, Iterable, List, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Callable, + ContextManager, + Iterable, + List, + Optional, + Tuple, + Union, +) from dvc.exceptions import FileMissingError from dvc.exceptions import IsADirectoryError as DvcIsADirectoryError @@ -463,6 +472,11 @@ class Repo: return brancher(self, *args, **kwargs) + def switch(self, rev: str) -> ContextManager[str]: + from dvc.repo.brancher import switch + + return switch(self, rev) + def used_objs( # noqa: PLR0913 self, targets=None, diff --git a/dvc/repo/brancher.py b/dvc/repo/brancher.py index 7a33312ff..93768edfd 100644 --- a/dvc/repo/brancher.py +++ b/dvc/repo/brancher.py @@ -1,8 +1,15 @@ import logging -from typing import Optional +from contextlib import contextmanager +from typing import TYPE_CHECKING, Iterator, Optional, Tuple +from scmrepo.git import Git + +from dvc.exceptions import NotDvcRepoError from dvc.scm import iter_revs +if TYPE_CHECKING: + from dvc.repo import Repo + logger = logging.getLogger(__name__) @@ -83,23 +90,68 @@ def brancher( # noqa: E302 ) try: - from dvc.fs import GitFileSystem - for sha, names in found_revs.items(): - logger.trace( # type: ignore[attr-defined] - "switching fs to revision %s [%s]", sha[:7], ", ".join(names) - ) - self.fs = GitFileSystem(scm=scm, rev=sha) - self.root_dir = self.fs.path.join("/", *repo_root_parts) - - if cwd_parts: - cwd = self.fs.path.join("/", *cwd_parts) # type: ignore[unreachable] - self.fs.path.chdir(cwd) - - # ignore revs that don't contain repo root - # (i.e. revs from before a subdir=True repo was init'ed) - if self.fs.exists(self.root_dir): + try: + _switch_fs(self, sha, repo_root_parts, cwd_parts) yield sha if sha_only else ",".join(names) + except NotDvcRepoError: + # ignore revs that don't contain repo root + # (i.e. revs from before a subdir=True repo was init'ed) + pass finally: self.fs = saved_fs self.root_dir = saved_root + + +def _switch_fs( + repo: "Repo", + rev: str, + repo_root_parts: Tuple[str, ...], + cwd_parts: Tuple[str, ...], +): + from dvc.fs import GitFileSystem, LocalFileSystem + + if rev == "workspace": + logger.trace("switching fs to workspace") # type: ignore[attr-defined] + repo.fs = LocalFileSystem(url=repo.root_dir) + return + + logger.trace("switching fs to revision %s", rev[:7]) # type: ignore[attr-defined] + assert isinstance(repo.scm, Git) + fs = GitFileSystem(scm=repo.scm, rev=rev) + root_dir = repo.fs.path.join("/", *repo_root_parts) + if not fs.exists(root_dir): + raise NotDvcRepoError(f"Commit '{rev[:7]}' does not contain a DVC repo") + + repo.fs = fs + repo.root_dir = root_dir + + if cwd_parts: + cwd = repo.fs.path.join("/", *cwd_parts) + repo.fs.path.chdir(cwd) + + +@contextmanager +def switch(repo: "Repo", rev: str) -> Iterator[str]: + """Switch to a specific revision.""" + from dvc.scm import resolve_rev + + if rev != "workspace": + rev = resolve_rev(repo.scm, rev) + + repo_root_parts = () + if repo.fs.path.isin(repo.root_dir, repo.scm.root_dir): + repo_root_parts = repo.fs.path.relparts(repo.root_dir, repo.scm.root_dir) + + cwd_parts = () + if repo.fs.path.isin(repo.fs.path.getcwd(), repo.scm.root_dir): + cwd_parts = repo.fs.path.relparts(repo.fs.path.getcwd(), repo.scm.root_dir) + + saved_fs = repo.fs + saved_root = repo.root_dir + try: + _switch_fs(repo, rev, repo_root_parts, cwd_parts) + yield rev + finally: + repo.fs = saved_fs + repo.root_dir = saved_root diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index 8f3b2037b..e9db9428e 100644 --- a/dvc/repo/experiments/utils.py +++ b/dvc/repo/experiments/utils.py @@ -35,6 +35,7 @@ from .refs import ( if TYPE_CHECKING: from dvc.repo import Repo + from dvc.scm import NoSCM EXEC_TMP_DIR = "exps" @@ -108,7 +109,9 @@ def exp_refs_by_rev(scm: "Git", rev: str) -> Generator[ExpRefInfo, None, None]: def exp_refs_by_baseline( - scm: "Git", revs: Optional[Set[str]] = None, url: Optional[str] = None + scm: "Git", + revs: Optional[Set[str]] = None, + url: Optional[str] = None, ) -> Mapping[str, List[ExpRefInfo]]: """Iterate over all experiment refs with the specified baseline.""" all_exp_refs = exp_refs(scm, url) @@ -233,7 +236,7 @@ def remove_exp_refs(scm: "Git", ref_infos: Iterable[ExpRefInfo]): scm.remove_ref(str(ref_info)) -def fix_exp_head(scm: "Git", ref: Optional[str]) -> Optional[str]: +def fix_exp_head(scm: Union["Git", "NoSCM"], ref: Optional[str]) -> Optional[str]: if ref: name, tail = Git.split_ref_pattern(ref) if name == "HEAD" and scm.get_ref(EXEC_BASELINE): diff --git a/dvc/scm.py b/dvc/scm.py index 30e05ec73..fd3239522 100644 --- a/dvc/scm.py +++ b/dvc/scm.py @@ -165,7 +165,7 @@ def clone(url: str, to_path: str, **kwargs): raise CloneError("SCM error") from exc -def resolve_rev(scm: "Git", rev: str) -> str: +def resolve_rev(scm: Union["Git", "NoSCM"], rev: str) -> str: from scmrepo.exceptions import RevError as InternalRevError from dvc.repo.experiments.utils import fix_exp_head @@ -173,6 +173,7 @@ def resolve_rev(scm: "Git", rev: str) -> str: try: return scm.resolve_rev(fix_exp_head(scm, rev)) except InternalRevError as exc: + assert isinstance(scm, Git) # `scm` will only resolve git branch and tag names, # if rev is not a sha it may be an abbreviated experiment name if not (rev == "HEAD" or rev.startswith("refs/")):