index: storage: don't forget to merge storage info

src/dvc_data/index/index.py

@@ -300,14 +300,32 @@ def __delitem__(self, key):
         del self._map[key]
 
     def __getitem__(self, key):
+        storages = []
         for prefix, storage in self._map.items():
             if len(prefix) > len(key):
                 continue
 
             if key[: len(prefix)] == prefix:
-                return storage
-
-        raise StorageKeyError(key)
+                storages.append((prefix, storage))
+
+        if not storages:
+            raise StorageKeyError(key)
+
+        storages = sorted(storages, key=lambda entry: len(entry[0]), reverse=True)
+        data = None
+        cache = None
+        remote = None
+        for _, storage in storages:
+            if data is None:
+                data = storage.data
+            if cache is None:
+                cache = storage.cache
+            if remote is None:
+                remote = storage.remote
+            if data and cache and remote:
+                break
+
+        return StorageInfo(data=data, cache=cache, remote=remote)
 
     def __iter__(self):
         yield from self._map.keys()

tests/index/test_storage.py

@@ -0,0 +1,49 @@
+from dvc_data.index import FileStorage, ObjectStorage, StorageInfo, StorageMapping
+
+
+def test_map_get(tmp_upath, as_filesystem, odb):
+    smap = StorageMapping()
+
+    data = FileStorage(key=(), fs=as_filesystem(tmp_upath.fs), path=str(tmp_upath))
+    cache = FileStorage(
+        key=("dir",), fs=as_filesystem(tmp_upath.fs), path=str(tmp_upath)
+    )
+    remote = FileStorage(
+        key=("dir", "subdir"), fs=as_filesystem(tmp_upath.fs), path=str(tmp_upath)
+    )
+    foo_cache = ObjectStorage(key=("dir", "foo"), odb=odb)
+
+    smap[()] = StorageInfo(data=data)
+    smap[("dir",)] = StorageInfo(cache=cache)
+    smap[("dir", "subdir")] = StorageInfo(remote=remote)
+    smap[("dir", "foo")] = StorageInfo(cache=foo_cache)
+
+    sinfo = smap[()]
+    assert sinfo.data == data
+    assert sinfo.cache is None
+    assert sinfo.remote is None
+
+    sinfo = smap[("dir",)]
+    assert sinfo.data == data
+    assert sinfo.cache == cache
+    assert sinfo.remote is None
+
+    sinfo = smap[("dir", "foo")]
+    assert sinfo.data == data
+    assert sinfo.cache == foo_cache
+    assert sinfo.remote is None
+
+    sinfo = smap[("dir", "subdir")]
+    assert sinfo.data == data
+    assert sinfo.cache == cache
+    assert sinfo.remote == remote
+
+    sinfo = smap[("dir", "subdir", "file")]
+    assert sinfo.data == data
+    assert sinfo.cache == cache
+    assert sinfo.remote == remote
+
+    sinfo = smap[("dir", "subdir", "subsubdir", "otherfile")]
+    assert sinfo.data == data
+    assert sinfo.cache == cache
+    assert sinfo.remote == remote