fix: coerce doc_ids in update/remove like get

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
 

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):