diff --git a/python/tests/test_collection.py b/python/tests/test_collection.py index b03d1e3cf..49fa91550 100644 --- a/python/tests/test_collection.py +++ b/python/tests/test_collection.py @@ -426,6 +426,26 @@ def test_collection_optimize(self, test_collection: Collection): # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionFetch: + @pytest.mark.parametrize("ids", [None, 1, b"doc_1", "", ["doc_1", None], [""]]) + def test_collection_fetch_rejects_invalid_ids(self, ids): + collection = Collection.__new__(Collection) + collection._obj = MagicMock() + + with pytest.raises(ValueError, match="ids must"): + collection.fetch(ids) + + collection._obj.Fetch.assert_not_called() + + def test_collection_fetch_accepts_tuple_ids(self): + collection = Collection.__new__(Collection) + collection._obj = MagicMock() + collection._obj.Fetch.return_value = {} + + result = collection.fetch(("doc_1", "doc_2")) + + assert result == {} + collection._obj.Fetch.assert_called_once_with(["doc_1", "doc_2"], None, True) + def test_collection_fetch( self, collection_with_single_doc: Collection, single_doc: Doc ): @@ -807,6 +827,26 @@ def test_collection_upsert_batch( # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionDelete: + @pytest.mark.parametrize("ids", [None, 1, b"doc_1", "", ["doc_1", None], [""]]) + def test_collection_delete_rejects_invalid_ids(self, ids): + collection = Collection.__new__(Collection) + collection._obj = MagicMock() + + with pytest.raises(ValueError, match="ids must"): + collection.delete(ids) + + collection._obj.Delete.assert_not_called() + + def test_collection_delete_accepts_tuple_ids(self): + collection = Collection.__new__(Collection) + collection._obj = MagicMock() + collection._obj.Delete.return_value = ["ok_1", "ok_2"] + + result = collection.delete(("doc_1", "doc_2")) + + assert result == ["ok_1", "ok_2"] + collection._obj.Delete.assert_called_once_with(["doc_1", "doc_2"]) + def test_empty_collection_delete(self, test_collection: Collection, single_doc): result = test_collection.delete(single_doc.id) assert bool(result) diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py index a002d0de5..2dff1c613 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -14,6 +14,7 @@ from __future__ import annotations import warnings +from collections.abc import Sequence from typing import Optional, Union, overload from zvec._zvec import _Collection @@ -48,6 +49,24 @@ def _require_positive_integer(value, name: str) -> None: raise ValueError(f"{name} must be a positive integer") +DocIds = Union[str, list[str], tuple[str, ...]] + + +def _normalize_doc_ids(ids: DocIds) -> tuple[list[str], bool]: + if isinstance(ids, str): + if not ids: + raise ValueError("ids must contain non-empty strings") + return [ids], True + + if not isinstance(ids, Sequence) or isinstance(ids, (bytes, bytearray)): + raise ValueError("ids must be a string or a sequence of strings") + + id_list = ids if isinstance(ids, list) else list(ids) + if any(not isinstance(doc_id, str) or not doc_id for doc_id in id_list): + raise ValueError("ids must contain non-empty strings") + return id_list, False + + class Collection: """Represents an opened collection in Zvec. @@ -324,21 +343,20 @@ def delete(self, ids: str) -> Status: pass @overload - def delete(self, ids: list[str]) -> list[Status]: + def delete(self, ids: Union[list[str], tuple[str, ...]]) -> list[Status]: pass - def delete(self, ids: Union[str, list[str]]) -> Union[Status, list[Status]]: + def delete(self, ids: DocIds) -> Union[Status, list[Status]]: """Delete documents by ID. Args: - ids (Union[str, list[str]]): One or more document IDs to delete. + ids (Union[str, list[str], tuple[str, ...]]): One or more document IDs to delete. Returns: Union[Status, list[Status]]: If a single id was given, returns its Status; if a list was given, returns a list of Status objects. """ - is_single = isinstance(ids, str) - id_list = [ids] if isinstance(ids, str) else ids + id_list, is_single = _normalize_doc_ids(ids) results = self._obj.Delete(id_list) return results[0] if is_single else results @@ -353,7 +371,7 @@ def delete_by_filter(self, filter: str) -> None: # ========== Collection DQL-fetch Methods ========== def fetch( self, - ids: Union[str, list[str]], + ids: DocIds, *, output_fields: Optional[list[str]] = None, include_vector: bool = True, @@ -361,7 +379,7 @@ def fetch( """Retrieve documents by ID. Args: - ids (Union[str, list[str]]): Document IDs to fetch. + ids (Union[str, list[str], tuple[str, ...]]): Document IDs to fetch. output_fields (Optional[list[str]], optional): Scalar fields to include. If None, all fields are returned. Defaults to None. include_vector (bool, optional): Whether to include vector data in @@ -370,8 +388,8 @@ def fetch( Returns: dict[str, Doc]: Mapping from ID to document. Missing IDs are omitted. """ - ids = [ids] if isinstance(ids, str) else ids - docs = self._obj.Fetch(ids, output_fields, include_vector) + id_list, _ = _normalize_doc_ids(ids) + docs = self._obj.Fetch(id_list, output_fields, include_vector) return { doc_id: py_doc for doc_id, core_doc in docs.items()