brancher: add brancher.switch helper

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,

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

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):

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/")):