odb: add exists_prefix method

Author: skshetry

src/dvc_objects/db.py

@@ -47,6 +47,14 @@ def __hash__(self):
     def exists(self, oid: str) -> bool:
         return self.fs.exists(self.oid_to_path(oid))
 
+    def exists_prefix(self, short_oid: str) -> str:
+        ret = [oid for oid in self.all() if oid.startswith(short_oid)]
+        if not ret:
+            raise KeyError(short_oid)
+        if len(ret) == 1:
+            return ret[0]
+        raise ValueError(short_oid, ret)
+
     def move(self, from_info, to_info):
         self.fs.move(from_info, to_info)
 

tests/test_odb.py

@@ -54,6 +54,24 @@ def test_exists(memfs):
     assert odb.exists("1234")
 
 
+def test_exists_prefix(memfs):
+    odb = ObjectDB(memfs, "/odb")
+    with pytest.raises(KeyError):
+        assert odb.exists_prefix("12")
+
+    odb.add_bytes("123456", b"content")
+    assert odb.exists_prefix("123") == "123456"
+
+
+def test_exists_prefix_ambiguous(memfs):
+    odb = ObjectDB(memfs, "/odb")
+    odb.add_bytes("123456", b"content")
+    odb.add_bytes("123450", b"content")
+    with pytest.raises(ValueError) as exc:
+        assert odb.exists_prefix("12")
+    assert exc.value.args == ("12", ["123450", "123456"])
+
+
 def test_move(memfs):
     odb = ObjectDB(memfs, "/")
     odb.add_bytes("1234", b"content")