diff --git a/tests/test_tinydb.py b/tests/test_tinydb.py index 762a460b..1caece56 100644 --- a/tests/test_tinydb.py +++ b/tests/test_tinydb.py @@ -278,6 +278,26 @@ def transform(el): assert db.count(where('int') == 1) == 2 +def test_update_custom_transform_callable(db: TinyDB): + from collections.abc import MutableMapping + + from tinydb.operations import delete as delete_field + + doc_id = db.insert({'name': 'John', 'temp': True}) + + def remove_temp(doc: MutableMapping) -> None: + del doc['temp'] + + db.update(remove_temp, where('name') == 'John') + + assert db.get(doc_id=doc_id) == {'name': 'John'} + + db.update({'temp': True}, where('name') == 'John') + db.update(delete_field('temp'), where('name') == 'John') + + assert db.get(doc_id=doc_id) == {'name': 'John'} + + def test_update_ids(db: TinyDB): db.update({'int': 2}, doc_ids=[1, 2]) diff --git a/tinydb/table.py b/tinydb/table.py index 17a0fb60..bffd7eab 100644 --- a/tinydb/table.py +++ b/tinydb/table.py @@ -3,7 +3,7 @@ data in TinyDB. """ -from collections.abc import Callable, Iterable, Iterator, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping, MutableMapping from typing import ( NoReturn, Optional, @@ -403,7 +403,7 @@ def contains( def update( self, - fields: Union[Mapping, Callable[[Mapping], None]], + fields: Union[Mapping, Callable[[MutableMapping], None]], cond: Optional[QueryLike] = None, doc_ids: Optional[Iterable[int]] = None, ) -> list[int]: @@ -415,7 +415,7 @@ def update( the IDs that were actually updated. :param fields: the fields that the matching documents will have - or a method that will update the documents + or a callable that mutates matching documents in place :param cond: which documents to update :param doc_ids: a list of document IDs :returns: a list containing the updated document's ID @@ -508,7 +508,7 @@ def updater(table: dict): def update_multiple( self, updates: Iterable[ - tuple[Union[Mapping, Callable[[Mapping], None]], QueryLike] + tuple[Union[Mapping, Callable[[MutableMapping], None]], QueryLike] ], ) -> list[int]: """