From 95ae018ed8971ecd978cc4ea8e35e19a59db1bd6 Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Tue, 13 Sep 2022 17:25:52 +0900
Subject: [PATCH] push/fetch: support dirs in worktree remotes

---
 dvc/repo/fetch.py |  4 ++--
 dvc/repo/push.py  | 37 +++++++++++++++++++++++++++++--------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py
index 3183bb683..0f71a0239 100644
--- a/dvc/repo/fetch.py
+++ b/dvc/repo/fetch.py
@@ -34,8 +34,8 @@ def _fetch_worktree(repo, remote):
             if not out.is_in_repo:
                 continue
 
-            key = repo.fs.path.relparts(out.fs_path, repo.root_dir)
-            entry = repo.index.data["repo"][key]
+            workspace, key = out.index_key
+            entry = repo.index.data[workspace][key]
             out.hash_info = entry.hash_info
             out.meta = entry.meta
 
diff --git a/dvc/repo/push.py b/dvc/repo/push.py
index 2b4ee07c2..3e7a2d1ea 100644
--- a/dvc/repo/push.py
+++ b/dvc/repo/push.py
@@ -10,11 +10,14 @@ if TYPE_CHECKING:
 
 
 def _push_worktree(repo, remote):
-    from dvc_data.index import build, checkout
+    from dvc_data.hashfile.tree import tree_from_index
+    from dvc_data.index import build, checkout, collect
 
     index = repo.index.data["repo"]
-    checkout(index, remote.path, remote.fs)
+    checkout(index, remote.path, remote.fs, force=True)
     build(index, remote.path, remote.fs)
+    if any(out.isdir() for out in repo.index.outs):
+        collect(index, remote.path, remote.fs)
 
     for stage in repo.index.stages:
         for out in stage.outs:
@@ -24,12 +27,30 @@ def _push_worktree(repo, remote):
             if not out.is_in_repo:
                 continue
 
-            key = repo.fs.path.relparts(out.fs_path, repo.root_dir)
-            entry = repo.index.data["repo"][key]
-            out.hash_info = entry.hash_info
-            out.meta = entry.meta
-
-        stage.dvcfile.dump(stage)
+            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 = tree_from_index(index, key)
+                new_tree.digest()
+                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)
 
     return len(index)