checkout: use index checkout

Needed for https://github.com/iterative/dvc/pull/9424

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):
+def _build_out_changes(index, changes):
-    return sum(stage.get_all_files_number(filter_info) for stage, filter_info in pairs)
+    from dvc_data.index.checkout import MODIFY
 
+    out_keys = []
+    for out in index.outs:
+        if not out.use_cache:
+            continue
 
-def _collect_pairs(
+        ws, key = out.index_key
-    self: "Repo", targets, with_deps: bool, recursive: bool
+        if ws != "repo":
-) -> Set["StageInfo"]:
+            continue
-    from dvc.stage.exceptions import StageFileBadNameError, StageFileDoesNotExistError
 
-    pairs: Set["StageInfo"] = set()
+        out_keys.append(key)
-    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
 
-    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)
+    def onerror(target, exc):
-    total = get_all_files_numbers(pairs)
+        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)
+        try:
-        for stage, filter_info in pairs:
+            changes = icheckout(
-            result = stage.checkout(
+                new,
-                force=force,
+                self.root_dir,
-                progress_callback=cb,
+                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():
+        except IndexCheckoutError as exc:
-                stats[key].extend(_fspath_dir(path) for path in items)
+            raise CheckoutError(exc.paths, {}) from exc
 
-    if stats.get("failed"):
+    out_changes = _build_out_changes(view, changes)
-        raise CheckoutError(stats["failed"], stats)
+
+    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

dvc/repo/data.py

@@ -73,10 +73,8 @@ def _diff(
     ):
         if (
             change.typ == UNCHANGED
-            and change.old
+            and (not change.old or not change.old.hash_info)
-            and change.new
+            and (not change.new or not change.new.hash_info)
-            and not change.old.hash_info
-            and not change.new.hash_info
         ):
             # NOTE: emulating previous behaviour
             continue

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)
 

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

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"]]:

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",

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

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 (
+from dvc.exceptions import CheckoutError, CheckoutErrorSuggestGit, NoOutputOrStageError
-    CheckoutError,
-    CheckoutErrorSuggestGit,
-    ConfirmRemoveError,
-    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))