From aabdddc2b790ade7233774d7e082ab4a60e078db Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Sat, 14 Jan 2023 03:18:59 +0200
Subject: [PATCH] index: introduce build_data_index

This is a generally useful function that builds a data index from repo index.

This is currently only used in cloud versioning, but it is going to be used in
`status`, `diff` and other data-centric operations
---
 dvc/repo/index.py    | 48 ++++++++++++++++++++++++++++++++++++++++++++
 dvc/repo/worktree.py | 23 ++-------------------
 2 files changed, 50 insertions(+), 21 deletions(-)

diff --git a/dvc/repo/index.py b/dvc/repo/index.py
index c44b51791..a51d20df1 100644
--- a/dvc/repo/index.py
+++ b/dvc/repo/index.py
@@ -34,6 +34,7 @@ if TYPE_CHECKING:
     from dvc_data.hashfile.hash_info import HashInfo
     from dvc_data.index import DataIndex, DataIndexKey, DataIndexView
     from dvc_objects.db import ObjectDB
+    from dvc_objects.fs.base import FileSystem
 
 
 logger = logging.getLogger(__name__)
@@ -447,3 +448,50 @@ class IndexView:
             else:
                 data[workspace] = DataIndex()
         return data
+
+
+def build_data_index(
+    index: Union["Index", "IndexView"],
+    path: str,
+    fs: "FileSystem",
+    workspace: Optional[str] = "repo",
+) -> "DataIndex":
+    from dvc_data.index import DataIndex, DataIndexEntry
+    from dvc_data.index.build import build_entries, build_entry
+
+    data = DataIndex()
+    for out in index.outs:
+        if not out.use_cache:
+            continue
+
+        ws, key = out.index_key
+        if ws != workspace:
+            continue
+
+        parts = out.fs.path.relparts(out.fs_path, out.repo.root_dir)
+        out_path = fs.path.join(path, *parts)
+
+        try:
+            entry = build_entry(out_path, fs)
+        except FileNotFoundError:
+            entry = DataIndexEntry(path=out_path, fs=fs)
+
+        entry.key = key
+
+        if not entry.meta or not entry.meta.isdir:
+            data.add(entry)
+            continue
+
+        entry.loaded = True
+        data.add(entry)
+
+        for entry in build_entries(out_path, fs):
+            if not entry.key or entry.key == ("",):
+                # NOTE: whether the root will be returned by build_entries
+                # depends on the filesystem (e.g. local doesn't, but s3 does).
+                entry.key = key
+            else:
+                entry.key = key + entry.key
+            data.add(entry)
+
+    return data
diff --git a/dvc/repo/worktree.py b/dvc/repo/worktree.py
index d964b2718..f2c032642 100644
--- a/dvc/repo/worktree.py
+++ b/dvc/repo/worktree.py
@@ -97,6 +97,7 @@ def push_worktree(
     jobs: Optional[int] = None,
     **kwargs,
 ) -> int:
+    from dvc.repo.index import build_data_index
     from dvc_data.index import checkout
 
     view = worktree_view(
@@ -108,7 +109,7 @@ def push_worktree(
     new_index = view.data["repo"]
     if remote.worktree:
         logger.debug("indexing latest worktree for '%s'", remote.path)
-        old_index = _build_worktree_index(repo, remote, view)
+        old_index = build_data_index(view, remote.path, remote.fs)
         logger.debug("Pushing worktree changes to '%s'", remote.path)
     else:
         old_index = None
@@ -147,26 +148,6 @@ def push_worktree(
     return pushed
 
 
-def _build_worktree_index(
-    repo: "Repo", remote: "Remote", view: "IndexView"
-) -> "DataIndex":
-    from dvc_data.index import DataIndex
-    from dvc_data.index.build import build_entries
-
-    index = DataIndex()
-    for out in view.outs:
-        _workspace, key = out.index_key
-        parts = out.fs.path.relparts(out.fs_path, repo.root_dir)
-        path = remote.fs.path.join(remote.path, *parts)
-        for entry in build_entries(path, remote.fs):
-            if not entry.key or entry.key == ("",):
-                entry.key = key
-            else:
-                entry.key = key + entry.key
-            index.add(entry)
-    return index
-
-
 def _update_out_meta(
     out: "Output", index: Union["DataIndex", "DataIndexView"]
 ):