diff --git a/dvc/repo/checkout.py b/dvc/repo/checkout.py index 51c3dc5f7..77bc328c6 100644 --- a/dvc/repo/checkout.py +++ b/dvc/repo/checkout.py @@ -1,16 +1,12 @@ import logging import os -from typing import TYPE_CHECKING, Dict, List, Set +from typing import Dict, List from dvc.exceptions import CheckoutError, CheckoutErrorSuggestGit, NoOutputOrStageError from dvc.utils import relpath from . import locked -if TYPE_CHECKING: - from . import Repo - from .stage import StageInfo - logger = logging.getLogger(__name__) @@ -30,37 +26,41 @@ def _remove_unused_links(repo): return ret -def get_all_files_numbers(pairs): - return sum(stage.get_all_files_number(filter_info) for stage, filter_info in pairs) +def _build_out_changes(index, changes): + from dvc_data.index.checkout import MODIFY + out_keys = [] + for out in index.outs: + if not out.use_cache: + continue -def _collect_pairs( - self: "Repo", targets, with_deps: bool, recursive: bool -) -> Set["StageInfo"]: - from dvc.stage.exceptions import StageFileBadNameError, StageFileDoesNotExistError + ws, key = out.index_key + if ws != "repo": + continue - pairs: Set["StageInfo"] = set() - for target in targets: - try: - pairs.update( - self.stage.collect_granular( - target, with_deps=with_deps, recursive=recursive - ) - ) - except ( - StageFileDoesNotExistError, - StageFileBadNameError, - NoOutputOrStageError, - ) as exc: - if not target: - raise - raise CheckoutErrorSuggestGit(target) from exc + out_keys.append(key) - return pairs + out_changes = {} + for objects in changes.values(): + for change in objects: + for out_key in out_keys: + if ( + len(out_key) > len(change.key) + or change.key[: len(out_key)] != out_key + ): + continue + + if change.key == out_key: + out_changes[out_key] = change.typ + elif not out_changes.get(out_key): + out_changes[out_key] = MODIFY + break + + return out_changes @locked -def checkout( +def checkout( # noqa: C901 self, targets=None, with_deps=False, @@ -70,13 +70,18 @@ def checkout( allow_missing=False, **kwargs, ): + from dvc import prompt from dvc.fs.callbacks import Callback + from dvc.repo.index import build_data_index + from dvc.stage.exceptions import StageFileBadNameError, StageFileDoesNotExistError + from dvc_data.hashfile.checkout import CheckoutError as IndexCheckoutError + from dvc_data.index.checkout import ADD, DELETE, MODIFY + from dvc_data.index.checkout import checkout as icheckout stats: Dict[str, List[str]] = { "added": [], "deleted": [], "modified": [], - "failed": [], } if not targets: targets = [None] @@ -85,28 +90,78 @@ def checkout( if isinstance(targets, str): targets = [targets] - pairs = _collect_pairs(self, targets, with_deps, recursive) - total = get_all_files_numbers(pairs) + def onerror(target, exc): + if target and isinstance( + exc, + ( + StageFileDoesNotExistError, + StageFileBadNameError, + NoOutputOrStageError, + ), + ): + raise CheckoutErrorSuggestGit(target) from exc + raise # pylint: disable=misplaced-bare-raise + + view = self.index.targets_view( + targets, + recursive=recursive, + with_deps=with_deps, + onerror=onerror, + ) + + with Callback.as_tqdm_callback( + unit="entry", + desc="Building data index", + ) as cb: + old = build_data_index( + view, self.root_dir, self.fs, compute_hash=True, callback=cb + ) + + new = view.data["repo"] + + for out in view.outs: + out.changed_cache() + with Callback.as_tqdm_callback( unit="file", desc="Checkout", - disable=total == 0, ) as cb: - cb.set_size(total) - for stage, filter_info in pairs: - result = stage.checkout( - force=force, - progress_callback=cb, + try: + changes = icheckout( + new, + self.root_dir, + self.fs, + old=old, + callback=cb, + delete=True, + prompt=prompt.confirm, + update_meta=False, relink=relink, - filter_info=filter_info, + force=force, allow_missing=allow_missing, + state=self.state, **kwargs, ) - for key, items in result.items(): - stats[key].extend(_fspath_dir(path) for path in items) + except IndexCheckoutError as exc: + raise CheckoutError(exc.paths, {}) from exc - if stats.get("failed"): - raise CheckoutError(stats["failed"], stats) + out_changes = _build_out_changes(view, changes) + + typ_map = {ADD: "added", DELETE: "deleted", MODIFY: "modified"} + for key, typ in out_changes.items(): + out_path = self.fs.path.join(self.root_dir, *key) + self.state.save_link(out_path, self.fs) + stats[typ_map[typ]].append(_fspath_dir(out_path)) + + failed = [] + for out in view.outs: + if not out.use_cache: + continue + + if not out.hash_info: + failed.append(_fspath_dir(out.fs_path)) + + if not allow_missing and failed: + raise CheckoutError(failed, stats) - del stats["failed"] return stats diff --git a/dvc/repo/data.py b/dvc/repo/data.py index 4bc80de50..209d3b293 100644 --- a/dvc/repo/data.py +++ b/dvc/repo/data.py @@ -73,10 +73,8 @@ def _diff( ): if ( change.typ == UNCHANGED - and change.old - and change.new - and not change.old.hash_info - and not change.new.hash_info + and (not change.old or not change.old.hash_info) + and (not change.new or not change.new.hash_info) ): # NOTE: emulating previous behaviour continue diff --git a/dvc/repo/imports.py b/dvc/repo/imports.py index d6f27a905..87154cca0 100644 --- a/dvc/repo/imports.py +++ b/dvc/repo/imports.py @@ -110,7 +110,7 @@ def save_imports( desc="Downloading imports from source", unit="files", ) as cb: - checkout(data_view, tmpdir, cache.fs, callback=cb, storage="data") + checkout(data_view, tmpdir, cache.fs, callback=cb, storage="remote") md5(data_view) save(data_view, odb=cache, hardlink=True) diff --git a/dvc/repo/index.py b/dvc/repo/index.py index 1a3227c99..ead9b2ca4 100644 --- a/dvc/repo/index.py +++ b/dvc/repo/index.py @@ -19,6 +19,7 @@ from typing import ( from funcy.debug import format_time from dvc.fs import LocalFileSystem +from dvc.fs.callbacks import DEFAULT_CALLBACK from dvc.utils.objects import cached_property if TYPE_CHECKING: @@ -26,6 +27,7 @@ if TYPE_CHECKING: from pygtrie import Trie from dvc.dependency import Dependency, ParamsDependency + from dvc.fs.callbacks import Callback from dvc.output import Output from dvc.repo import Repo from dvc.repo.stage import StageInfo @@ -113,14 +115,18 @@ def collect_files( def _load_data_from_outs(index, prefix, outs): - from dvc_data.index import DataIndexEntry + from dvc_data.index import DataIndexEntry, Meta + parents = set() for out in outs: if not out.use_cache: continue ws, key = out.index_key + for key_len in range(1, len(key)): + parents.add((ws, key[:key_len])) + entry = DataIndexEntry( key=key, meta=out.meta, @@ -137,6 +143,11 @@ def _load_data_from_outs(index, prefix, outs): # index.view() work correctly. index[(*prefix, ws, *key)] = entry + for ws, key in parents: + index[(*prefix, ws, *key)] = DataIndexEntry( + key=key, meta=Meta(isdir=True), loaded=True + ) + def _load_storage_from_out(storage_map, key, out): from dvc.config import NoRemoteError @@ -163,7 +174,7 @@ def _load_storage_from_out(storage_map, key, out): if out.stage.is_import: dep = out.stage.deps[0] - storage_map.add_data(FileStorage(key, dep.fs, dep.fs_path)) + storage_map.add_remote(FileStorage(key, dep.fs, dep.fs_path)) class Index: @@ -534,10 +545,12 @@ class IndexView: lambda: _DataPrefixes(set(), set()) ) for out, filter_info in self._filtered_outs: + if not out.use_cache: + continue workspace, key = out.index_key if filter_info and out.fs.path.isin(filter_info, out.fs_path): key = key + out.fs.path.relparts(filter_info, out.fs_path) - entry = self._index.data[workspace][key] + entry = self._index.data[workspace].get(key) if entry and entry.meta and entry.meta.isdir: prefixes[workspace].recursive.add(key) prefixes[workspace].explicit.update(key[:i] for i in range(len(key), 0, -1)) @@ -550,6 +563,9 @@ class IndexView: ret: Dict[str, Set["DataIndexKey"]] = defaultdict(set) for out, filter_info in self._filtered_outs: + if not out.use_cache: + continue + workspace, key = out.index_key if filter_info and out.fs.path.isin(filter_info, out.fs_path): key = key + out.fs.path.relparts(filter_info, out.fs_path) @@ -579,14 +595,15 @@ class IndexView: return data -def build_data_index( +def build_data_index( # noqa: C901 index: Union["Index", "IndexView"], path: str, fs: "FileSystem", workspace: str = "repo", compute_hash: Optional[bool] = False, + callback: "Callback" = DEFAULT_CALLBACK, ) -> "DataIndex": - from dvc_data.index import DataIndex, DataIndexEntry + from dvc_data.index import DataIndex, DataIndexEntry, Meta from dvc_data.index.build import build_entries, build_entry from dvc_data.index.save import build_tree @@ -595,9 +612,16 @@ def build_data_index( ignore = index.repo.dvcignore data = DataIndex() + parents = set() for key in index.data_keys.get(workspace, set()): out_path = fs.path.join(path, *key) + for key_len in range(1, len(key)): + parents.add(key[:key_len]) + + if not fs.exists(out_path): + continue + try: out_entry = build_entry( out_path, @@ -606,10 +630,11 @@ def build_data_index( state=index.repo.state, ) except FileNotFoundError: - out_entry = DataIndexEntry() + continue out_entry.key = key data.add(out_entry) + callback.relative_update(1) if not out_entry.meta or not out_entry.meta.isdir: continue @@ -628,6 +653,7 @@ def build_data_index( entry.key = key + entry.key data.add(entry) + callback.relative_update(1) if compute_hash: tree_meta, tree = build_tree(data, key) @@ -635,5 +661,14 @@ def build_data_index( out_entry.hash_info = tree.hash_info out_entry.loaded = True data.add(out_entry) + callback.relative_update(1) + + for key in parents: + parent_path = fs.path.join(path, *key) + if not fs.exists(parent_path): + continue + direntry = DataIndexEntry(key=key, meta=Meta(isdir=True), loaded=True) + data.add(direntry) + callback.relative_update(1) return data diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 7acb1ac10..26c7b34b9 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -717,11 +717,6 @@ class Stage(params.StageParams): for out in self.outs ) - def get_all_files_number(self, filter_info=None) -> int: - return sum( - out.get_files_number(filter_info) for out in self.filter_outs(filter_info) - ) - def get_used_objs( self, *args, **kwargs ) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]: diff --git a/pyproject.toml b/pyproject.toml index c17eb75d1..5cbbce1a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ "configobj>=5.0.6", "distro>=1.3", "dpath<3,>=2.1.0", - "dvc-data>=0.53.0,<0.54", + "dvc-data>=0.54.3,<0.55", "dvc-http>=2.29.0", "dvc-render>=0.3.1,<1", "dvc-studio-client>=0.9.2,<1", diff --git a/tests/func/test_add.py b/tests/func/test_add.py index 0ab55d62e..d19e39bf3 100644 --- a/tests/func/test_add.py +++ b/tests/func/test_add.py @@ -336,11 +336,11 @@ def test_should_update_state_entry_for_file_after_add(mocker, dvc, tmp_dir): os.rename("foo", "foo.back") ret = main(["checkout"]) assert ret == 0 - assert file_md5_counter.mock.call_count == 1 + assert file_md5_counter.mock.call_count == 2 ret = main(["status"]) assert ret == 0 - assert file_md5_counter.mock.call_count == 1 + assert file_md5_counter.mock.call_count == 2 def test_should_update_state_entry_for_directory_after_add(mocker, dvc, tmp_dir): @@ -362,11 +362,11 @@ def test_should_update_state_entry_for_directory_after_add(mocker, dvc, tmp_dir) os.rename("data", "data.back") ret = main(["checkout"]) assert ret == 0 - assert file_md5_counter.mock.call_count == 5 + assert file_md5_counter.mock.call_count == 6 ret = main(["status"]) assert ret == 0 - assert file_md5_counter.mock.call_count == 6 + assert file_md5_counter.mock.call_count == 7 def test_add_commit(tmp_dir, dvc): diff --git a/tests/func/test_checkout.py b/tests/func/test_checkout.py index 29d5dd8d4..6d5beab26 100644 --- a/tests/func/test_checkout.py +++ b/tests/func/test_checkout.py @@ -8,12 +8,7 @@ import pytest from dvc.cli import main from dvc.dvcfile import PROJECT_FILE, load_file -from dvc.exceptions import ( - CheckoutError, - CheckoutErrorSuggestGit, - ConfirmRemoveError, - NoOutputOrStageError, -) +from dvc.exceptions import CheckoutError, CheckoutErrorSuggestGit, NoOutputOrStageError from dvc.fs import LocalFileSystem, system from dvc.stage import Stage from dvc.stage.exceptions import StageFileDoesNotExistError @@ -82,7 +77,7 @@ def test_checkout_corrupted_cache_file(tmp_dir, dvc): fd.write("1") with pytest.raises(CheckoutError): - dvc.checkout(force=True) + dvc.checkout(force=True, relink=True) assert not os.path.isfile("foo") assert not os.path.isfile(cache) @@ -114,7 +109,7 @@ def test_checkout_corrupted_cache_dir(tmp_dir, dvc): fobj.write("1") with pytest.raises(CheckoutError): - dvc.checkout(force=True) + dvc.checkout(force=True, relink=True) assert not os.path.exists(cache) @@ -294,19 +289,6 @@ def test_checkout_directory(tmp_dir, dvc): assert os.path.exists("data") -def test_checkout_hook(mocker, tmp_dir, dvc): - """Test that dvc checkout handles EOFError gracefully, which is what - it will experience when running in a git hook. - """ - tmp_dir.dvc_gen({"data": {"foo": "foo"}}) - mocker.patch("sys.stdout.isatty", return_value=True) - mocker.patch("dvc.prompt.input", side_effect=EOFError) - - (tmp_dir / "data").gen("test", "test") - with pytest.raises(ConfirmRemoveError): - dvc.checkout() - - def test_checkout_suggest_git(tmp_dir, dvc, scm): with pytest.raises(CheckoutErrorSuggestGit) as e: dvc.checkout(targets="gitbranch") @@ -636,27 +618,6 @@ def test_checkout_recursive(tmp_dir, dvc): } -def test_checkout_for_external_outputs(tmp_dir, dvc, workspace): - workspace.gen("foo", "foo") - - file_path = workspace / "foo" - dvc.add("remote://workspace/foo") - - odb = dvc.cloud.get_remote_odb("workspace") - odb.fs.remove(str(file_path)) - assert not file_path.exists() - - stats = dvc.checkout(force=True) - assert stats == {**empty_checkout, "added": ["remote://workspace/foo"]} - assert file_path.exists() - - workspace.gen("foo", "foo\nfoo") - - stats = dvc.checkout(force=True) - assert stats == {**empty_checkout, "modified": ["remote://workspace/foo"]} - assert file_path.read_text() == "foo" - - def test_checkouts_with_different_addressing(tmp_dir, dvc, run_copy): tmp_dir.gen({"foo": "foo", "lorem": "lorem"}) run_copy("foo", "bar", name="copy-foo-bar") @@ -725,22 +686,6 @@ def test_checkouts_for_pipeline_tracked_outs(tmp_dir, dvc, scm, run_copy): assert set(dvc.checkout()["added"]) == {"bar", "ipsum"} -def test_checkout_external_modified_file(tmp_dir, dvc, scm, mocker, workspace): - # regression: happened when file in external output changed and checkout - # was attempted without force, dvc checks if it's present in its cache - # before asking user to remove it. - workspace.gen("foo", "foo") - dvc.add("remote://workspace/foo", external=True) - scm.add(["foo.dvc"]) - scm.commit("add foo") - - workspace.gen("foo", "foobar") # messing up the external outputs - mocker.patch("dvc.prompt.confirm", return_value=True) - dvc.checkout() - - assert (workspace / "foo").read_text() == "foo" - - def test_checkout_executable(tmp_dir, dvc): tmp_dir.dvc_gen("foo", "foo") @@ -817,7 +762,7 @@ def test_checkout_partial_unchanged(tmp_dir, dvc): assert len(stats["modified"]) == 1 stats = dvc.checkout(str(data_dir / "empty_sub_dir")) - assert not any(stats.values()) + assert stats == {**empty_checkout, "modified": ["data" + os.sep]} dvc.checkout(str(data_dir))