diff --git a/tests/test_tinydb.py b/tests/test_tinydb.py index 9ff866fe..ca2356f8 100644 --- a/tests/test_tinydb.py +++ b/tests/test_tinydb.py @@ -364,6 +364,18 @@ def test_doc_id_missing_consistency(db: TinyDB): assert db.remove(doc_ids=[99]) == [] +def test_doc_ids_string_coercion(db: TinyDB): + # ``get(doc_id='1')`` already works via str coercion; ``update`` / + # ``remove`` must accept the same string IDs (see #639). + assert db.get(doc_id='1') is not None + assert db.update({'int': 9}, doc_ids=['1']) == [1] + doc = db.get(doc_id=1) + assert doc is not None + assert doc['int'] == 9 + assert db.remove(doc_ids=['1']) == [1] + assert db.get(doc_id=1) is None + + def test_update_multiple(db: TinyDB): assert len(db) == 3 diff --git a/tinydb/table.py b/tinydb/table.py index f4027bfa..cf9edd8d 100644 --- a/tinydb/table.py +++ b/tinydb/table.py @@ -309,12 +309,14 @@ def get( @overload def get( - self, cond: Optional[QueryLike], doc_id: int, doc_ids: Optional[list] = ... + self, cond: Optional[QueryLike], doc_id: Union[int, str], + doc_ids: Optional[list] = ... ) -> Optional[Document]: ... @overload def get( - self, *, cond: Optional[QueryLike] = ..., doc_id: int, doc_ids: Optional[list] = ..., + self, *, cond: Optional[QueryLike] = ..., doc_id: Union[int, str], + doc_ids: Optional[list] = ..., ) -> Optional[Document]: ... @overload @@ -335,7 +337,7 @@ def get( def get( self, cond: Optional[QueryLike] = None, - doc_id: Optional[int] = None, + doc_id: Optional[Union[int, str]] = None, doc_ids: Optional[list] = None ): """ @@ -360,8 +362,9 @@ def get( if raw_doc is None: return None - # Convert the raw data to the document class - return self.document_class(raw_doc, doc_id) + # Convert the raw data to the document class. Coerce the ID the same + # way other read paths do so str IDs like '1' become document_id_class. + return self.document_class(raw_doc, self.document_id_class(doc_id)) elif doc_ids is not None: # Filter the table by extracting out all those documents which @@ -423,7 +426,7 @@ def update( self, fields: Union[Mapping, Callable[[MutableMapping], None]], cond: Optional[QueryLike] = None, - doc_ids: Optional[Iterable[int]] = None, + doc_ids: Optional[Iterable[Union[int, str]]] = None, ) -> list[int]: """ Update all matching documents to have a given set of fields. @@ -457,7 +460,9 @@ def perform_update(table, doc_id): # (see issue #591). The list of *actually* updated IDs is # determined inside the updater so it reflects the table state at # write time. - requested_ids = list(doc_ids) + # Coerce IDs with document_id_class so string IDs like '1' match + # table keys the same way get(doc_id='1') does (see issue #639). + requested_ids = [self.document_id_class(doc_id) for doc_id in doc_ids] updated_ids: list[int] = [] def updater(table: dict): @@ -620,7 +625,7 @@ def upsert(self, document: Mapping, cond: Optional[QueryLike] = None) -> list[in def remove( self, cond: Optional[QueryLike] = None, - doc_ids: Optional[Iterable[int]] = None, + doc_ids: Optional[Iterable[Union[int, str]]] = None, ) -> list[int]: """ Remove all matching documents. @@ -640,7 +645,9 @@ def remove( # ``get(doc_ids=...)`` (see issue #591). The list of *actually* # removed IDs is determined inside the updater so it reflects the # table state at write time. - requested_ids = list(doc_ids) + # Coerce IDs with document_id_class so string IDs like '1' match + # table keys the same way get(doc_id='1') does (see issue #639). + requested_ids = [self.document_id_class(doc_id) for doc_id in doc_ids] removed_ids: list[int] = [] def updater(table: dict):