diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index 7f345983c..269d16bdd 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -17,20 +17,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def _fetch_worktree(repo, remote): - from dvc_data.index import save - - index = repo.index.data["repo"] - for key, entry in index.iteritems(): - entry.fs = remote.fs - entry.path = remote.fs.path.join( - remote.path, - *key, - ) - save(index) - return len(index) - - @locked def fetch( self, @@ -59,6 +45,7 @@ def fetch( remote is configured """ from dvc.repo.imports import save_imports + from dvc.repo.worktree import fetch as fetch_worktree from dvc_data.hashfile.transfer import TransferResult if isinstance(targets, str): @@ -67,7 +54,7 @@ def fetch( with suppress(NoRemoteError): _remote = self.cloud.get_remote(name=remote) if _remote.worktree: - return _fetch_worktree(self, _remote) + return fetch_worktree(self, _remote) failed_count = 0 diff --git a/dvc/repo/push.py b/dvc/repo/push.py index bb4369107..c0ef6d7d1 100644 --- a/dvc/repo/push.py +++ b/dvc/repo/push.py @@ -9,48 +9,6 @@ if TYPE_CHECKING: from dvc_objects.db.base import ObjectDB -def _push_worktree(repo, remote): - from dvc_data.index import checkout - from dvc_data.index.save import build_tree - - index = repo.index.data["repo"] - checkout(index, remote.path, remote.fs) - - for stage in repo.index.stages: - for out in stage.outs: - if not out.use_cache: - continue - - if not out.is_in_repo: - continue - - workspace, key = out.index_key - index = repo.index.data[workspace] - entry = index[key] - if out.isdir(): - old_tree = out.get_obj() - entry.hash_info = old_tree.hash_info - entry.meta = out.meta - for subkey, entry in index.iteritems(key): - if entry.meta.isdir: - continue - fs_path = repo.fs.path.join(repo.root_dir, *subkey) - _, hash_info = old_tree.get( - repo.fs.path.relparts(fs_path, out.fs_path) - ) - entry.hash_info = hash_info - tree_meta, new_tree = build_tree(index, key) - out.obj = new_tree - out.hash_info = new_tree.hash_info - out.meta = tree_meta - else: - out.hash_info = entry.hash_info - out.meta = entry.meta - stage.dvcfile.dump(stage, with_files=True, update_pipeline=False) - - return len(index) - - @locked def push( self, @@ -68,9 +26,11 @@ def push( odb: Optional["ObjectDB"] = None, include_imports=False, ): + from dvc.repo.worktree import push as push_worktree + _remote = self.cloud.get_remote(name=remote) if _remote.worktree: - return _push_worktree(self, _remote) + return push_worktree(self, _remote) used_run_cache = ( self.stage_cache.push(remote, odb=odb) if run_cache else [] diff --git a/dvc/repo/worktree.py b/dvc/repo/worktree.py new file mode 100644 index 000000000..10f93f5a7 --- /dev/null +++ b/dvc/repo/worktree.py @@ -0,0 +1,61 @@ +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from dvc.cloud import Remote + from dvc.repo import Repo + + +def fetch(repo: "Repo", remote: "Remote") -> int: + from dvc_data.index import save + + index = repo.index.data["repo"] + for key, entry in index.iteritems(): + entry.fs = remote.fs + entry.path = remote.fs.path.join( + remote.path, + *key, + ) + save(index) + return len(index) + + +def push(repo: "Repo", remote: "Remote") -> int: + from dvc_data.index import checkout + from dvc_data.index.save import build_tree + + index = repo.index.data["repo"] + checkout(index, remote.path, remote.fs) + + for stage in repo.index.stages: + for out in stage.outs: + if not out.use_cache: + continue + + if not out.is_in_repo: + continue + + workspace, key = out.index_key + index = repo.index.data[workspace] + entry = index[key] + if out.isdir(): + old_tree = out.get_obj() + entry.hash_info = old_tree.hash_info + entry.meta = out.meta + for subkey, entry in index.iteritems(key): + if entry.meta.isdir: + continue + fs_path = repo.fs.path.join(repo.root_dir, *subkey) + _, hash_info = old_tree.get( + repo.fs.path.relparts(fs_path, out.fs_path) + ) + entry.hash_info = hash_info + tree_meta, new_tree = build_tree(index, key) + out.obj = new_tree + out.hash_info = new_tree.hash_info + out.meta = tree_meta + else: + out.hash_info = entry.hash_info + out.meta = entry.meta + stage.dvcfile.dump(stage, with_files=True, update_pipeline=False) + + return len(index)