Revert "scm: use dulwich backend when fetching exps during clone/pull" (#9041)

Revert "scm: use dulwich backend when fetching exps during clone/pull (#9023)"

This reverts commit ed1ab7fa75.

dvc/cachemgr.py

@@ -9,7 +9,6 @@ def _get_odb(repo, settings, fs=None):
         return None
 
     cls, config, fs_path = get_cloud_fs(repo, **settings)
-    config["tmp_dir"] = repo.tmp_dir
     fs = fs or cls(**config)
     return get_odb(fs, fs_path, state=repo.state, **config)
 

dvc/cli/__init__.py

@@ -39,12 +39,17 @@ def parse_args(argv=None):
 
 def _log_unknown_exceptions() -> None:
     from dvc.info import get_dvc_info
-    from dvc.logger import FOOTER
+    from dvc.ui import ui
+    from dvc.utils import colorize
 
     logger.exception("unexpected error")
     if logger.isEnabledFor(logging.DEBUG):
         logger.debug("Version info for developers:\n%s", get_dvc_info())
-    logger.info(FOOTER)
+
+    q = colorize("Having any troubles?", "yellow")
+    link = colorize("https://dvc.org/support", "blue")
+    footer = f"\n{q} Hit us up at {link}, we are always happy to help!"
+    ui.error_write(footer)
 
 
 def _log_exceptions(exc: Exception) -> Optional[int]:

dvc/commands/plots.py

@@ -150,7 +150,7 @@ class CmdPlots(CmdBase):
                 render_html(
                     renderers=renderers,
                     output_file=output_file,
-                    template_path=self._html_template_path(),
+                    html_template=self._html_template_path(),
                 )
 
                 ui.write(output_file.as_uri())

dvc/commands/update.py

@@ -31,8 +31,7 @@ class CmdUpdate(CmdBase):
 def add_parser(subparsers, parent_parser):
     UPDATE_HELP = (
         "Update data artifact imported (via dvc import or dvc import-url) "
-        "from an external DVC repository or URL, or update data from a "
-        "DVC worktree remote."
+        "from an external DVC repository or URL."
     )
     update_parser = subparsers.add_parser(
         "update",

dvc/config_schema.py

@@ -51,6 +51,8 @@ def ByUrl(mapping):
         # Windows absolute paths should really have scheme == "" (local)
         if os.name == "nt" and len(parsed.scheme) == 1 and parsed.netloc == "":
             return schemas[""](data)
+        if parsed.netloc == "":
+            return schemas[""](data)
         if parsed.scheme not in schemas:
             raise Invalid(f"Unsupported URL type {parsed.scheme}://")
 
@@ -262,6 +264,7 @@ SCHEMA = {
     # section for experimental features
     "feature": {
         Optional("machine", default=False): Bool,
+        Optional("data_index_cache", default=False): Bool,
         # enabled by default. It's of no use, kept for backward compatibility.
         Optional("parametrization", default=True): Bool,
     },

dvc/data_cloud.py

@@ -18,10 +18,11 @@ logger = logging.getLogger(__name__)
 
 
 class Remote:
-    def __init__(self, name: str, path: str, fs: "FileSystem", **config):
+    def __init__(self, name: str, path: str, fs: "FileSystem", *, index=None, **config):
         self.path = path
         self.fs = fs
         self.name = name
+        self.index = index
 
         self.worktree: bool = config.pop("worktree", False)
         self.config = config
@@ -75,7 +76,11 @@ class DataCloud:
 
             fs = cls(**config)
             config["tmp_dir"] = self.repo.index_db_dir
-            return Remote(name, fs_path, fs, **config)
+            if self.repo.data_index is not None:
+                index = self.repo.data_index.view(("remotes", name))
+            else:
+                index = None
+            return Remote(name, fs_path, fs, index=index, **config)
 
         if bool(self.repo.config["remote"]):
             error_msg = (

dvc/logger.py

@@ -9,16 +9,6 @@ import colorama
 
 from dvc.progress import Tqdm
 
-FOOTER = (
-    "\n{yellow}Having any troubles?{nc}"
-    " Hit us up at {blue}https://dvc.org/support{nc},"
-    " we are always happy to help!"
-).format(
-    blue=colorama.Fore.BLUE,
-    nc=colorama.Fore.RESET,
-    yellow=colorama.Fore.YELLOW,
-)
-
 
 def addLoggingLevel(levelName, levelNum, methodName=None):
     """

dvc/render/convert.py

@@ -29,10 +29,8 @@ def _group_by_rev(datapoints):
 
 
 def to_json(renderer, split: bool = False) -> List[Dict]:
-    from copy import deepcopy
-
     if renderer.TYPE == "vega":
-        grouped = _group_by_rev(deepcopy(renderer.datapoints))
+        grouped = _group_by_rev(renderer.datapoints)
         if split:
             content = renderer.get_filled_template(skip_anchors=["data"])
         else:

dvc/render/converter/vega.py

@@ -1,4 +1,3 @@
-from copy import deepcopy
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from funcy import first, last
@@ -233,7 +232,7 @@ class VegaConverter(Converter):
         for i, (y_file, y_field) in enumerate(ys):
             if num_xs > 1:
                 x_file, x_field = xs[i]
-            datapoints = deepcopy(file2datapoints.get(y_file, []))
+            datapoints = [d.copy() for d in file2datapoints.get(y_file, [])]
 
             if props_update.get("y", None) == "dvc_inferred_y_value":
                 _update_from_field(

dvc/repo/__init__.py

@@ -23,6 +23,7 @@ if TYPE_CHECKING:
     from dvc.stage import Stage
     from dvc.types import DictStrAny
     from dvc_data.hashfile.state import StateBase
+    from dvc_data.index import DataIndex
 
     from .experiments import Experiments
     from .index import Index
@@ -179,6 +180,7 @@ class Repo:
         self._fs_conf = {"repo_factory": repo_factory}
         self._fs = fs or localfs
         self._scm = scm
+        self._data_index = None
 
         if rev and not fs:
             self._scm = scm = SCM(root_dir or os.curdir)
@@ -341,6 +343,31 @@ class Repo:
         # fs.
         self._reset()
 
+    @property
+    def data_index(self) -> Optional["DataIndex"]:
+        from appdirs import user_cache_dir
+        from fsspec.utils import tokenize
+
+        from dvc_data.index import DataIndex
+
+        if not self.config["feature"].get("data_index_cache"):
+            return None
+
+        if self._data_index is None:
+            cache_dir = user_cache_dir(self.config.APPNAME, self.config.APPAUTHOR)
+            index_dir = os.path.join(
+                cache_dir,
+                "index",
+                "data",
+                # scm.root_dir and repo.root_dir don't match for subrepos
+                tokenize((self.scm.root_dir, self.root_dir)),
+            )
+            os.makedirs(index_dir, exist_ok=True)
+
+            self._data_index = DataIndex.open(os.path.join(index_dir, "db.db"))
+
+        return self._data_index
+
     def __repr__(self):
         return f"{self.__class__.__name__}: '{self.root_dir}'"
 
@@ -553,18 +580,17 @@ class Repo:
     def close(self):
         self.scm.close()
         self.state.close()
+        if self._data_index is not None:
+            self._data_index.close()
 
-    def _reset_cached_indecies(self):
+    def _reset(self):
+        self.scm._reset()  # pylint: disable=protected-access
+        self.state.close()
         self.__dict__.pop("index", None)
         self.__dict__.pop("dvcignore", None)
         self.__dict__.pop("dvcfs", None)
         self.__dict__.pop("datafs", None)
 
-    def _reset(self):
-        self.state.close()
-        self.scm._reset()  # pylint: disable=protected-access
-        self._reset_cached_indecies()
-
     def __enter__(self):
         return self
 

dvc/repo/brancher.py

@@ -1,7 +1,10 @@
+import logging
 from typing import Optional
 
 from dvc.scm import iter_revs
 
+logger = logging.getLogger(__name__)
+
 
 def brancher(  # noqa: E302
     self,
@@ -60,6 +63,7 @@ def brancher(  # noqa: E302
 
     scm = self.scm
 
+    logger.trace("switching fs to workspace")  # type: ignore[attr-defined]
     self.fs = LocalFileSystem(url=self.root_dir)
     yield "workspace"
 
@@ -82,7 +86,9 @@ def brancher(  # noqa: E302
         from dvc.fs import GitFileSystem
 
         for sha, names in found_revs.items():
-            self._reset_cached_indecies()  # pylint: disable=W0212
+            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)
 

dvc/repo/data.py

@@ -13,17 +13,20 @@ from typing import (
     Union,
 )
 
+from dvc.ui import ui
+
 if TYPE_CHECKING:
     from dvc.repo import Repo
     from dvc.scm import Git, NoSCM
     from dvc_data.index import DataIndex
+    from dvc_data.index.diff import Change
 
 
 def posixpath_to_os_path(path: str) -> str:
     return path.replace(posixpath.sep, os.path.sep)
 
 
-def _adapt_typ(typ):
+def _adapt_typ(typ: str) -> str:
     from dvc_data.index.diff import ADD, DELETE, MODIFY
 
     if typ == MODIFY:
@@ -38,7 +41,7 @@ def _adapt_typ(typ):
     return typ
 
 
-def _adapt_path(change):
+def _adapt_path(change: "Change") -> str:
     isdir = False
     if change.new and change.new.meta:
         isdir = change.new.meta.isdir
@@ -54,8 +57,8 @@ def _diff(
     old: "DataIndex",
     new: "DataIndex",
     *,
-    granular: Optional[bool] = False,
-    with_missing: Optional[bool] = False,
+    granular: bool = False,
+    with_missing: bool = False,
 ) -> Dict[str, List[str]]:
     from dvc_data.index.diff import UNCHANGED, UNKNOWN, diff
 
@@ -94,7 +97,7 @@ def _diff(
             with_missing
             and change.old
             and change.old.hash_info
-            and not old.storage_map[change.key].cache.exists(change.old.hash_info.value)
+            and not old.storage_map.cache_exists(change.old)
         ):
             # NOTE: emulating previous behaviour
             _add_change("not_in_cache", change)
@@ -143,14 +146,18 @@ def _git_info(scm: Union["Git", "NoSCM"], untracked_files: str = "all") -> GitIn
 def _diff_index_to_wtree(repo: "Repo", **kwargs: Any) -> Dict[str, List[str]]:
     from .index import build_data_index
 
-    workspace = build_data_index(repo.index, repo.root_dir, repo.fs, compute_hash=True)
+    with ui.status("Building workspace index"):
+        workspace = build_data_index(
+            repo.index, repo.root_dir, repo.fs, compute_hash=True
+        )
 
-    return _diff(
-        repo.index.data["repo"],
-        workspace,
-        with_missing=True,
-        **kwargs,
-    )
+    with ui.status("Calculating diff between index/workspace"):
+        return _diff(
+            repo.index.data["repo"],
+            workspace,
+            with_missing=True,
+            **kwargs,
+        )
 
 
 def _diff_head_to_index(
@@ -164,7 +171,8 @@ def _diff_head_to_index(
 
         head_index = repo.index.data["repo"]
 
-    return _diff(head_index, index, **kwargs)
+    with ui.status("Calculating diff between head/index"):
+        return _diff(head_index, index, **kwargs)
 
 
 class Status(TypedDict):

dvc/repo/diff.py

@@ -1,39 +1,103 @@
 import logging
 import os
 from collections import defaultdict
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    Iterable,
-    Iterator,
-    List,
-    Optional,
-    Set,
-    Tuple,
-    cast,
-)
+from typing import Dict, List, Optional
 
 from dvc.exceptions import PathMissingError
 from dvc.repo import locked
-from dvc.utils.collections import ensure_list
-
-if TYPE_CHECKING:
-    from dvc.fs import DVCFileSystem, FileSystem
-    from dvc.output import Output
-    from dvc.repo import Repo
-    from dvc_data.hashfile.obj import HashFile
-    from dvc_data.hashfile.tree import Tree
+from dvc.ui import ui
 
 logger = logging.getLogger(__name__)
 
 
+def _path(entry):
+    if entry and entry.meta and entry.meta.isdir:
+        return os.path.join(*entry.key, "")
+    return os.path.join(*entry.key)
+
+
+def _hash(entry):
+    if entry and entry.hash_info:
+        return entry.hash_info.value
+    return None
+
+
+def _diff(old, new, with_missing=False):
+    from dvc_data.index.diff import ADD, DELETE, MODIFY, RENAME
+    from dvc_data.index.diff import diff as idiff
+
+    ret: "Dict[str, List[Dict]]" = {
+        "added": [],
+        "deleted": [],
+        "modified": [],
+        "renamed": [],
+        "not in cache": [],
+    }
+
+    for change in idiff(
+        old,
+        new,
+        with_renames=True,
+        hash_only=True,
+    ):
+        if change.typ == ADD:
+            ret["added"].append(
+                {
+                    "path": _path(change.new),
+                    "hash": _hash(change.new),
+                }
+            )
+        elif change.typ == DELETE:
+            ret["deleted"].append(
+                {
+                    "path": _path(change.old),
+                    "hash": _hash(change.old),
+                }
+            )
+        elif change.typ == MODIFY:
+            ret["modified"].append(
+                {
+                    "path": _path(change.old),
+                    "hash": {
+                        "old": _hash(change.old),
+                        "new": _hash(change.new),
+                    },
+                }
+            )
+        elif change.typ == RENAME:
+            ret["renamed"].append(
+                {
+                    "path": {
+                        "old": _path(change.old),
+                        "new": _path(change.new),
+                    },
+                    "hash": _hash(change.old),
+                }
+            )
+
+        if (
+            with_missing
+            and change.old
+            and change.old.hash_info
+            and not old.storage_map.cache_exists(change.old)
+        ):
+            ret["not in cache"].append(
+                {
+                    "path": _path(change.old),
+                    "hash": _hash(change.old),
+                }
+            )
+
+    return ret if any(ret.values()) else {}
+
+
 @locked
 def diff(
     self,
     a_rev: str = "HEAD",
     b_rev: Optional[str] = None,
     targets: Optional[List[str]] = None,
+    recursive: bool = False,
 ):
     """
     By default, it compares the workspace with the last commit's fs.
@@ -42,231 +106,58 @@ def diff(
     the concept of `index`, but it keeps the same interface, thus,
     `dvc diff` would be the same as `dvc diff HEAD`.
     """
-
     if self.scm.no_commits:
         return {}
 
-    from dvc.fs.dvc import DVCFileSystem
+    indexes = {}
+    missing_targets = defaultdict(set)
+    with_missing = False
+    if not b_rev:
+        b_rev = "workspace"
+        with_missing = True
 
-    dvcfs = DVCFileSystem(repo=self)
-    targets = [dvcfs.from_os_path(target) for target in ensure_list(targets)]
-
-    b_rev = b_rev if b_rev else "workspace"
-    results = {}
-    missing_targets = {}
     for rev in self.brancher(revs=[a_rev, b_rev]):
-        if rev == "workspace" and rev != b_rev:
+        if rev == "workspace" and b_rev != "workspace":
             # brancher always returns workspace, but we only need to compute
             # workspace paths/checksums if b_rev was None
             continue
 
-        targets_paths = None
-        if targets:
-            # convert targets to paths, and capture any missing targets
-            targets_paths, missing_targets[rev] = _targets_to_paths(dvcfs, targets)
+        def onerror(target, _exc):
+            # pylint: disable-next=cell-var-from-loop
+            missing_targets[rev].add(target)  # noqa: B023
 
-        results[rev] = _paths_checksums(self, targets_paths)
-
-    if targets:
-        # check for overlapping missing targets between a_rev and b_rev
-        for target in set(missing_targets[a_rev]) & set(missing_targets[b_rev]):
-            raise PathMissingError(target, self)
-
-    old = results[a_rev]
-    new = results[b_rev]
-
-    # Compare paths between the old and new fs.
-    # set() efficiently converts dict keys to a set
-    added = sorted(set(new) - set(old))
-    deleted_or_missing = set(old) - set(new)
-    if b_rev == "workspace":
-        # missing status is only applicable when diffing local workspace
-        # against a commit
-        missing = sorted(_filter_missing(dvcfs, deleted_or_missing))
-    else:
-        missing = []
-    deleted = sorted(deleted_or_missing - set(missing))
-    modified = sorted(set(old) & set(new))
-
-    # Cases when file was changed and renamed are resulted
-    # in having deleted and added record
-    # To cover such cases we need to change hashing function
-    # to produce rolling/chunking hash
-
-    renamed = _calculate_renamed(new, old, added, deleted)
-
-    for renamed_item in renamed:
-        added.remove(renamed_item["path"]["new"])
-        deleted.remove(renamed_item["path"]["old"])
-
-    ret = {
-        "added": [{"path": path, "hash": new[path]} for path in added],
-        "deleted": [{"path": path, "hash": old[path]} for path in deleted],
-        "modified": [
-            {"path": path, "hash": {"old": old[path], "new": new[path]}}
-            for path in modified
-            if old[path] != new[path]
-        ],
-        "renamed": renamed,
-        "not in cache": [{"path": path, "hash": old[path]} for path in missing],
-    }
-
-    return ret if any(ret.values()) else {}
-
-
-def _paths_checksums(
-    repo: "Repo", targets: Optional[List[str]]
-) -> Dict[str, Optional[str]]:
-    """
-    A dictionary of checksums addressed by relpaths collected from
-    the current fs outputs.
-
-    To help distinguish between a directory and a file output,
-    the former one will come with a trailing slash in the path:
-
-        directory: "data/"
-        file:      "data"
-    """
-
-    return dict(_output_paths(repo, targets))
-
-
-def _output_paths(  # noqa: C901
-    repo: "Repo", targets: Optional[List[str]]
-) -> Iterable[Tuple[str, Optional[str]]]:
-    from dvc.fs import LocalFileSystem
-    from dvc_data.hashfile.build import build
-
-    on_working_fs = isinstance(repo.fs, LocalFileSystem)
-
-    def _exists(output: "Output") -> bool:
-        if on_working_fs:
-            return output.exists
-        return True
-
-    def _to_path(output: "Output") -> str:
-        relparts = output.fs.path.relparts(output.fs_path)
-        base = os.path.join(*relparts)
-        if output.is_dir_checksum:
-            return os.path.join(base, "")
-        return base
-
-    for output in repo.index.outs:
-        if not _exists(output):
-            continue
-
-        yield_output = targets is None or any(
-            output.fs.path.isin_or_eq(output.fs_path, target) for target in targets
+        view = self.index.targets_view(
+            targets,
+            onerror=onerror,
+            recursive=recursive,
         )
 
-        obj: Optional["HashFile"]
-        if on_working_fs:
-            _, _, obj = build(
-                repo.cache.repo,
-                output.fs_path,
-                repo.cache.repo.fs,
-                "md5",
-                dry_run=True,
-                ignore=output.dvcignore,
-            )
-            hash_info = obj.hash_info
-        else:
-            hash_info = output.hash_info
-            obj = output.get_obj()
+        if rev == "workspace":
+            from .index import build_data_index
 
-        if yield_output:
-            yield _to_path(output), hash_info.value
-
-        if not obj:
-            continue
-
-        if output.is_dir_checksum and (
-            yield_output
-            or (
-                isinstance(targets, list)
-                and any(
-                    output.fs.path.isin(target, output.fs_path) for target in targets
+            with ui.status("Building workspace index"):
+                data = build_data_index(
+                    view,
+                    self.root_dir,
+                    self.fs,
+                    compute_hash=True,
                 )
-            )
-        ):
-            obj = cast("Tree", obj)
-            yield from _dir_output_paths(output.fs, output.fs_path, obj, targets)
-
-
-def _dir_output_paths(
-    fs: "FileSystem",
-    fs_path: str,
-    obj: "Tree",
-    targets: Optional[Iterable[str]] = None,
-) -> Iterator[Tuple[str, str]]:
-    base = os.path.join(*fs.path.relparts(fs_path))
-    for key, _, oid in obj:
-        fname = fs.path.join(fs_path, *key)
-        if targets is None or any(
-            fs.path.isin_or_eq(fname, target) for target in targets
-        ):
-            # pylint: disable=no-member
-            yield os.path.join(base, *key), oid.value
-
-
-def _filter_missing(dvcfs: "DVCFileSystem", paths: Set[str]) -> Iterator[str]:
-    for path in paths:
-        fs_path = dvcfs.from_os_path(path)
-        try:
-            info = dvcfs.info(fs_path)
-            dvc_info = info.get("dvc_info", {})
-            entry = dvc_info.get("entry")
-            if (
-                entry
-                and info["type"] == "directory"
-                and not dvcfs.repo.cache.local.exists(entry.hash_info.value)
-            ):
-                yield path
-        except FileNotFoundError:
-            pass
-
-
-def _targets_to_paths(
-    dvcfs: "DVCFileSystem", targets: List[str]
-) -> Tuple[List[str], List[str]]:
-    paths: List[str] = []
-    missing: List[str] = []
-
-    if not targets:
-        return paths, missing
-
-    for target in targets:
-        if dvcfs.exists(target):
-            paths.append(dvcfs.repo.fs.path.abspath(target))
         else:
-            missing.append(target)
+            data = view.data["repo"]
 
-    return paths, missing
+        assert rev not in indexes
+        indexes[rev] = data
 
+    if targets:
+        old_missing = missing_targets.get(a_rev, set())
+        new_missing = missing_targets.get(b_rev, set())
 
-def _calculate_renamed(
-    new: Dict[str, Any],
-    old: Dict[str, Any],
-    added: Iterable[str],
-    deleted: Iterable[str],
-) -> List[Dict[str, Dict[str, str]]]:
-    old_inverted: Dict[str, List[str]] = defaultdict(list)
-    # It is needed to be dict of lists to cover cases
-    # when repo has paths with same hash
-    for path, path_hash in old.items():
-        old_inverted[path_hash].append(path)
+        # check for overlapping missing targets between a_rev and b_rev
+        for target in old_missing & new_missing:
+            raise PathMissingError(target, self)
 
-    renamed = []
-    for path in added:
-        path_hash = new[path]
-        old_paths = old_inverted[path_hash]
-        try:
-            iterator = enumerate(old_paths)
-            index = next(idx for idx, path in iterator if path in deleted)
-        except StopIteration:
-            continue
+    old = indexes[a_rev]
+    new = indexes[b_rev]
 
-        old_path = old_paths.pop(index)
-        renamed.append({"path": {"old": old_path, "new": path}, "hash": path_hash})
-
-    return renamed
+    with ui.status("Calculating diff"):
+        return _diff(old, new, with_missing=with_missing)

dvc/repo/experiments/__init__.py

@@ -96,19 +96,9 @@ class Experiments:
     def reproduce_one(
         self,
         tmp_dir: bool = False,
-        machine: Optional[str] = None,
         **kwargs,
     ):
         """Reproduce and checkout a single (standalone) experiment."""
-        if not (tmp_dir or machine):
-            staged, _, _ = self.scm.status(untracked_files="no")
-            if staged:
-                logger.warning(
-                    "Your workspace contains staged Git changes which will be "
-                    "unstaged before running this experiment."
-                )
-                self.scm.reset()
-
         exp_queue: "BaseStashQueue" = (
             self.tempdir_queue if tmp_dir else self.workspace_queue
         )

dvc/repo/experiments/queue/base.py

@@ -323,12 +323,12 @@ class BaseStashQueue(ABC):
         .. _Hydra Override:
             https://hydra.cc/docs/next/advanced/override_grammar/basic/
         """
-        with self.scm.detach_head(client="dvc") as orig_head:
-            stash_head = orig_head
-            if baseline_rev is None:
-                baseline_rev = orig_head
+        with self.scm.stash_workspace(reinstate_index=True) as workspace:
+            with self.scm.detach_head(client="dvc") as orig_head:
+                stash_head = orig_head
+                if baseline_rev is None:
+                    baseline_rev = orig_head
 
-            with self.scm.stash_workspace() as workspace:
                 try:
                     if workspace:
                         self.stash.apply(workspace)

dvc/repo/experiments/queue/celery.py

@@ -336,7 +336,8 @@ class LocalCeleryQueue(BaseStashQueue):
                         task_id,
                         remained_entries[entry],
                     )
-                    self.celery.backend.mark_as_failure(task_id, None)
+                    backend = self.celery.backend
+                    backend.mark_as_failure(task_id, None)  # type: ignore[attr-defined]
 
         if remained_revs:
             raise CannotKillTasksError(remained_revs)

dvc/repo/experiments/queue/remove.py

@@ -40,7 +40,7 @@ def remove_tasks(  # noqa: C901
             msg,
             queue_entry,
         ) in celery_queue._iter_queued():  # pylint: disable=protected-access
-            if queue_entry.stash_rev in stash_revs:
+            if queue_entry.stash_rev in stash_revs and msg.delivery_tag:
                 celery_queue.celery.reject(msg.delivery_tag)
     finally:
         celery_queue.stash.remove_revs(list(stash_revs.values()))
@@ -56,7 +56,8 @@ def remove_tasks(  # noqa: C901
             result: AsyncResult = AsyncResult(task_id)
             if result is not None:
                 result.forget()
-            celery_queue.celery.purge(msg.delivery_tag)
+            if msg.delivery_tag:
+                celery_queue.celery.purge(msg.delivery_tag)
     finally:
         if celery_queue.failed_stash:
             celery_queue.failed_stash.remove_revs(failed_stash_revs)

dvc/repo/imports.py

@@ -115,7 +115,7 @@ def save_imports(
                 desc="Downloading imports from source",
                 unit="files",
             ) as cb:
-                checkout(data_view, tmpdir, cache.fs, callback=cb)
+                checkout(data_view, tmpdir, cache.fs, callback=cb, storage="data")
             md5(data_view)
             save(data_view, odb=cache, hardlink=True)
 

dvc/repo/index.py

@@ -112,6 +112,60 @@ def collect_files(
         dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
 
 
+def _load_data_from_outs(index, prefix, outs):
+    from dvc_data.index import DataIndexEntry
+
+    for out in outs:
+        if not out.use_cache:
+            continue
+
+        ws, key = out.index_key
+
+        entry = DataIndexEntry(
+            key=key,
+            meta=out.meta,
+            hash_info=out.hash_info,
+        )
+
+        if out.stage.is_import and not out.stage.is_repo_import:
+            dep = out.stage.deps[0]
+            entry.meta = dep.meta
+            entry.hash_info = dep.hash_info
+
+        # FIXME PyGTrie-based DataIndex doesn't remove entry.key during
+        # index.add, so we have to set the entry manually here to make
+        # index.view() work correctly.
+        index[(*prefix, ws, *key)] = entry
+
+
+def _load_storage_from_out(storage_map, key, out):
+    from dvc.config import NoRemoteError
+    from dvc_data.index import FileStorage, ObjectStorage
+
+    if out.odb:
+        storage_map.add_data(ObjectStorage(key, out.odb))
+    storage_map.add_cache(ObjectStorage(key, out.cache))
+    try:
+        remote = out.repo.cloud.get_remote(out.remote)
+        if remote.fs.version_aware:
+            storage_map.add_remote(
+                FileStorage(
+                    key=key,
+                    fs=remote.fs,
+                    path=remote.fs.path.join(remote.path, *key),
+                    index=remote.index,
+                )
+            )
+        else:
+            storage_map.add_remote(ObjectStorage(key, remote.odb, index=remote.index))
+    except NoRemoteError:
+        pass
+
+    if out.stage.is_import and not out.stage.is_repo_import:
+        dep = out.stage.deps[0]
+        storage_map.add_data(FileStorage(key, dep.fs, dep.fs_path))
+
+
 class Index:
     def __init__(
         self,
@@ -282,50 +336,60 @@ class Index:
         return dict(by_workspace)
 
     @cached_property
-    def data(self) -> "Dict[str, DataIndex]":
-        from collections import defaultdict
-
-        from dvc.config import NoRemoteError
-        from dvc_data.index import DataIndex, DataIndexEntry, Storage
-
-        by_workspace: dict = defaultdict(DataIndex)
-
-        by_workspace["repo"] = DataIndex()
-        by_workspace["local"] = DataIndex()
+    def data_tree(self):
+        from dvc_data.hashfile import Tree
 
+        tree = Tree()
         for out in self.outs:
             if not out.use_cache:
                 continue
 
-            workspace, key = out.index_key
-            data_index = by_workspace[workspace]
+            ws, key = out.index_key
 
-            entry = DataIndexEntry(
-                key=key,
-                meta=out.meta,
-                hash_info=out.hash_info,
-            )
+            tree.add((ws, *key), out.meta, out.hash_info)
 
-            storage = Storage(odb=out.odb, cache=out.cache)
-            try:
-                storage.remote = self.repo.cloud.get_remote_odb(out.remote)
-            except NoRemoteError:
-                pass
+        tree.digest()
 
-            if out.stage.is_import and not out.stage.is_repo_import:
-                dep = out.stage.deps[0]
+        return tree
 
-                entry.meta = dep.meta
-                entry.hash_info = dep.hash_info
+    @cached_property
+    def data(self) -> "Dict[str, DataIndex]":
+        from dvc_data.index import DataIndex
 
-                storage.odb = dep.odb
-                storage.fs = dep.fs
-                storage.path = dep.fs_path
+        prefix: "DataIndexKey"
+        loaded = False
 
-            data_index.add(entry)
-            data_index.storage_map[key] = storage
+        index = self.repo.data_index
+        if index is None:
+            index = DataIndex()
 
-        return dict(by_workspace)
+        prefix = (self.data_tree.hash_info.value,)
+        if prefix in index.ls((), detail=False):
+            loaded = True
+
+        try:
+            if not loaded:
+                _load_data_from_outs(index, prefix, self.outs)
+                index.commit()
+
+            by_workspace = {}
+            by_workspace["repo"] = index.view((*prefix, "repo"))
+            by_workspace["local"] = index.view((*prefix, "local"))
+
+            for out in self.outs:
+                if not out.use_cache:
+                    continue
+
+                ws, key = out.index_key
+                if ws not in by_workspace:
+                    by_workspace[ws] = index.view((*prefix, ws))
+
+                data_index = by_workspace[ws]
+                _load_storage_from_out(data_index.storage_map, key, out)
+
+            return by_workspace
+        finally:
+            index.close()
 
     @staticmethod
     def _hash_targets(
@@ -526,6 +590,10 @@ def build_data_index(
     from dvc_data.index.build import build_entries, build_entry
     from dvc_data.index.save import build_tree
 
+    ignore = None
+    if workspace == "repo" and isinstance(fs, LocalFileSystem):
+        ignore = index.repo.dvcignore
+
     data = DataIndex()
     for key in index.data_keys.get(workspace, set()):
         out_path = fs.path.join(path, *key)
@@ -547,7 +615,11 @@ def build_data_index(
             continue
 
         for entry in build_entries(
-            out_path, fs, compute_hash=compute_hash, state=index.repo.state
+            out_path,
+            fs,
+            compute_hash=compute_hash,
+            state=index.repo.state,
+            ignore=ignore,
         ):
             if not entry.key or entry.key == ("",):
                 # NOTE: whether the root will be returned by build_entries

dvc/repo/worktree.py

@@ -119,12 +119,6 @@ def fetch_worktree(
     ):
         remote_obj = _get_remote(repo, remote_name, remote, "fetch")
         index = view.data["repo"]
-        for key in index.storage_map:
-            storage = index.storage_map[key]
-            storage.fs = remote_obj.fs
-            storage.path = remote_obj.fs.path.join(remote_obj.path, *key)
-            index.storage_map[key] = storage
-
         total = len(index)
         with Callback.as_tqdm_callback(
             unit="file",
@@ -132,7 +126,7 @@ def fetch_worktree(
             disable=total == 0,
         ) as cb:
             cb.set_size(total)
-            transferred += save(index, callback=cb, jobs=jobs)
+            transferred += save(index, callback=cb, jobs=jobs, storage="remote")
     return transferred
 
 
@@ -196,6 +190,11 @@ def push_worktree(
                     f"remote {remote_obj.name!r} does not support versioning"
                 ) from None
 
+        if remote_obj.index is not None:
+            for key, entry in new_index.iteritems():
+                remote_obj.index[key] = entry
+            remote_obj.index.commit()
+
         for out in view.outs:
             workspace, _key = out.index_key
             _merge_push_meta(out, repo.index.data[workspace], remote_obj.name)
@@ -357,6 +356,7 @@ def _fetch_out_changes(
             update_meta=False,
             meta_only=True,
             meta_cmp_key=partial(_meta_checksum, remote.fs),
+            storage="data",
             callback=cb,
         )
         out.save()

pyproject.toml

@@ -46,14 +46,14 @@ dependencies = [
     "pygtrie>=2.3.2",
     "dpath>=2.0.2,<3",
     "shtab>=1.3.4,<2",
-    "rich>=10.13.0",
+    "rich>=12.0.0",
     "pyparsing>=2.4.7",
     "typing-extensions>=3.7.4",
-    "scmrepo==0.1.9",
-    "dvc-render==0.1.1",
+    "scmrepo==0.1.11",
+    "dvc-render>=0.1.2",
     "dvc-task==0.1.11",
     "dvclive>=1.2.2",
-    "dvc-data==0.39.3",
+    "dvc-data==0.40.3",
     "dvc-http",
     "hydra-core>=1.1.0",
     "iterative-telemetry==0.0.7",
@@ -100,7 +100,7 @@ tests = [
     "filelock==3.9.0",
     "beautifulsoup4==4.11.2",
     "pywin32>=225; sys_platform == 'win32'",  # optional dependency
-    "pylint==2.16.1",
+    "pylint==2.16.2",
     # type-checking
     "mypy==0.991",
     "types-requests",
@@ -109,7 +109,7 @@ tests = [
     "types-appdirs",
     "types-colorama",
     "types-psutil",
-    "types-tqdm",
+    "types-tqdm<4.64.7.13",
 ]
 all = ["dvc[azure,gdrive,gs,hdfs,oss,s3,ssh,webdav,webhdfs]"]
 dev = ["dvc[azure,gdrive,gs,hdfs,oss,s3,ssh,webdav,webhdfs,tests]"]

tests/func/experiments/test_stash_exp.py

@@ -0,0 +1,67 @@
+import pytest
+from funcy import first
+
+from dvc.dependency.base import DependencyDoesNotExistError
+from dvc.exceptions import ReproductionError
+
+
+@pytest.mark.parametrize("tmp", [True, False])
+@pytest.mark.parametrize("staged", [True, False])
+def test_deleted(tmp_dir, scm, dvc, tmp, staged):
+    tmp_dir.scm_gen("file", "file", commit="commit file")
+    stage = dvc.stage.add(
+        cmd="cat file",
+        deps=["file"],
+        name="foo",
+    )
+    scm.add_commit(["dvc.yaml"], message="add dvc.yaml")
+
+    file = tmp_dir / "file"
+    file.unlink()
+    if staged:
+        scm.add(["file"])
+
+    with pytest.raises(ReproductionError) as exc_info:
+        dvc.experiments.run(stage.addressing, tmp_dir=tmp)
+
+    cause = exc_info._excinfo[1].__cause__
+    assert isinstance(cause, DependencyDoesNotExistError)
+    assert not file.exists()
+
+
+@pytest.mark.parametrize("tmp", [True, False])
+@pytest.mark.parametrize("staged", [True, False])
+def test_modified(tmp_dir, scm, dvc, caplog, tmp, staged):
+    tmp_dir.scm_gen("file", "file", commit="commit file")
+    stage = dvc.stage.add(
+        cmd="cat file",
+        name="foo",
+    )
+    scm.add_commit(["dvc.yaml"], message="add dvc.yaml")
+
+    (tmp_dir / "file").write_text("modified_file")
+    if staged:
+        scm.add(["file"])
+
+    results = dvc.experiments.run(stage.addressing, tmp_dir=tmp)
+
+    exp = first(results)
+    scm.checkout(exp, force=True)
+    assert (tmp_dir / "file").read_text() == "modified_file"
+
+
+@pytest.mark.parametrize("tmp", [True, False])
+def test_staged_new_file(tmp_dir, scm, dvc, tmp):
+    stage = dvc.stage.add(
+        cmd="cat file",
+        name="foo",
+    )
+    scm.add_commit(["dvc.yaml"], message="add dvc.yaml")
+
+    (tmp_dir / "file").write_text("file")
+    scm.add(["file"])
+
+    results = dvc.experiments.run(stage.addressing, tmp_dir=tmp)
+    exp = first(results)
+    fs = scm.get_fs(exp)
+    assert fs.exists("file")

tests/func/test_add.py

@@ -342,6 +342,13 @@ def test_double_add_unchanged_dir(tmp_dir, dvc):
     assert ret == 0
 
 
+@pytest.mark.skipif(os.name == "nt", reason="unsupported on Windows")
+def test_add_colon_in_filename(tmp_dir, dvc):
+    tmp_dir.gen("fo:o", "foo")
+    ret = main(["add", "fo:o"])
+    assert ret == 0
+
+
 def test_should_update_state_entry_for_file_after_add(mocker, dvc, tmp_dir):
     file_md5_counter = mocker.spy(dvc_data.hashfile.hash, "file_md5")
     tmp_dir.gen("foo", "foo")

tests/func/test_data_cloud.py

@@ -199,7 +199,8 @@ def test_verify_hashes(tmp_dir, scm, dvc, mocker, tmp_path_factory, local_remote
     hash_spy = mocker.spy(dvc_data.hashfile.hash, "file_md5")
 
     dvc.pull()
-    assert hash_spy.call_count == 0
+    # NOTE: 1 is for index.data_tree building
+    assert hash_spy.call_count == 1
 
     # Removing cache will invalidate existing state entries
     dvc.cache.local.clear()
@@ -207,7 +208,7 @@ def test_verify_hashes(tmp_dir, scm, dvc, mocker, tmp_path_factory, local_remote
     dvc.config["remote"]["upstream"]["verify"] = True
 
     dvc.pull()
-    assert hash_spy.call_count == 4
+    assert hash_spy.call_count == 6
 
 
 @flaky(max_runs=3, min_passes=1)

tests/func/test_diff.py

@@ -78,7 +78,12 @@ def test_no_cache_entry(tmp_dir, scm, dvc):
                 "hash": {"old": digest("first"), "new": digest("second")},
             }
         ],
-        "not in cache": [],
+        "not in cache": [
+            {
+                "path": "file",
+                "hash": digest("first"),
+            }
+        ],
         "renamed": [],
     }
 
@@ -265,7 +270,12 @@ def test_diff_no_cache(tmp_dir, scm, dvc):
     remove(str(tmp_dir / "dir"))
     diff = dvc.diff()
     assert diff["added"] == []
-    assert diff["deleted"] == []
+    assert diff["deleted"] == [
+        {
+            "path": os.path.join("dir", ""),
+            "hash": "f0f7a307d223921557c929f944bf5303.dir",
+        }
+    ]
     assert diff["renamed"] == []
     assert diff["modified"] == []
     assert diff["not in cache"] == [
@@ -460,8 +470,10 @@ def test_targets_single_dir_with_file(tmp_dir, scm, dvc):
         "renamed": [],
     }
 
-    assert dvc.diff(targets=["dir_with"]) == expected_result
-    assert dvc.diff(targets=["dir_with" + os.path.sep]) == expected_result
+    assert dvc.diff(targets=["dir_with"], recursive=True) == expected_result
+    assert (
+        dvc.diff(targets=["dir_with" + os.path.sep], recursive=True) == expected_result
+    )
 
 
 def test_targets_single_file_in_dir_with_file(tmp_dir, scm, dvc):

tests/unit/command/test_plots.py

@@ -285,7 +285,7 @@ def test_should_call_render(tmp_dir, mocker, capsys, plots_data, output):
     render_mock.assert_called_once_with(
         renderers=renderers,
         output_file=Path(tmp_dir / output / "index.html"),
-        template_path=None,
+        html_template=None,
     )