diff --git a/telegram/ext/_dispatcher.py b/telegram/ext/_dispatcher.py index 3bba0df15e0..f635cf32cf9 100644 --- a/telegram/ext/_dispatcher.py +++ b/telegram/ext/_dispatcher.py @@ -55,6 +55,7 @@ from telegram.ext._utils.stack import was_called_by if TYPE_CHECKING: + from telegram import Message from telegram.ext._jobqueue import Job from telegram.ext._builders import InitDispatcherBuilder @@ -680,6 +681,52 @@ def drop_user_data(self, user_id: int) -> None: if self.persistence: self.persistence.drop_user_data(user_id) + def migrate_chat_data( + self, message: 'Message' = None, old_chat_id: int = None, new_chat_id: int = None + ) -> None: + """Moves the contents of :attr:`chat_data` at key old_chat_id to the key new_chat_id. + Also updates the persistence by calling :attr:`update_persistence`. + + Warning: + + * Any data stored in :attr:`chat_data` at key `new_chat_id` will be overridden + * The key `old_chat_id` of :attr:`chat_data` will be deleted + + Args: + message (:class:`Message`, optional): A message with either + :attr:`telegram.Message.migrate_from_chat_id` or + :attr:`telegram.Message.migrate_to_chat_id`. + Mutually exclusive with passing ``old_chat_id`` and ``new_chat_id`` + + .. seealso: `telegram.ext.filters.StatusUpdate.MIGRATE` + old_chat_id (:obj:`int`, optional): The old chat ID. + Mutually exclusive with passing ``message`` + new_chat_id (:obj:`int`, optional): The new chat ID. + Mutually exclusive with passing ``message`` + + """ + if message and (old_chat_id or new_chat_id): + raise ValueError("Message and chat_id pair are mutually exclusive") + if not any((message, old_chat_id, new_chat_id)): + raise ValueError("chat_id pair or message must be passed") + + if message: + if message.migrate_from_chat_id is None and message.migrate_to_chat_id is None: + raise ValueError( + "Invalid message instance. The message must have either " + "`Message.migrate_from_chat_id` or `Message.migrate_to_chat_id`." + ) + + old_chat_id = message.migrate_from_chat_id or message.chat.id + new_chat_id = message.migrate_to_chat_id or message.chat.id + + elif not (isinstance(old_chat_id, int) and isinstance(new_chat_id, int)): + raise ValueError("old_chat_id and new_chat_id must be integers") + + self._chat_data[new_chat_id] = self._chat_data[old_chat_id] + self.drop_chat_data(old_chat_id) + self.update_persistence() + def update_persistence(self, update: object = None) -> None: """Update :attr:`user_data`, :attr:`chat_data` and :attr:`bot_data` in :attr:`persistence`. diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py index 6a077049a22..eecc123b484 100644 --- a/tests/test_dispatcher.py +++ b/tests/test_dispatcher.py @@ -745,6 +745,51 @@ def test_sensible_worker_thread_names(self, dp2): for thread_name in thread_names: assert thread_name.startswith(f"Bot:{dp2.bot.id}:worker:") + @pytest.mark.parametrize( + 'message', + [ + Message(message_id=1, chat=Chat(id=2, type=None), migrate_from_chat_id=1, date=None), + Message(message_id=1, chat=Chat(id=1, type=None), migrate_to_chat_id=2, date=None), + Message(message_id=1, chat=Chat(id=1, type=None), date=None), + None, + ], + ) + @pytest.mark.parametrize('old_chat_id', [None, 1, "1"]) + @pytest.mark.parametrize('new_chat_id', [None, 2, "1"]) + def test_migrate_chat_data(self, dp, message: 'Message', old_chat_id: int, new_chat_id: int): + def call(match: str): + with pytest.raises(ValueError, match=match): + dp.migrate_chat_data( + message=message, old_chat_id=old_chat_id, new_chat_id=new_chat_id + ) + + if message and (old_chat_id or new_chat_id): + call(r"^Message and chat_id pair are mutually exclusive$") + return + + if not any((message, old_chat_id, new_chat_id)): + call(r"^chat_id pair or message must be passed$") + return + + if message: + if message.migrate_from_chat_id is None and message.migrate_to_chat_id is None: + call(r"^Invalid message instance") + return + effective_old_chat_id = message.migrate_from_chat_id or message.chat.id + effective_new_chat_id = message.migrate_to_chat_id or message.chat.id + + elif not (isinstance(old_chat_id, int) and isinstance(new_chat_id, int)): + call(r"^old_chat_id and new_chat_id must be integers$") + return + else: + effective_old_chat_id = old_chat_id + effective_new_chat_id = new_chat_id + + dp.chat_data[effective_old_chat_id]['key'] = "test" + dp.migrate_chat_data(message=message, old_chat_id=old_chat_id, new_chat_id=new_chat_id) + assert effective_old_chat_id not in dp.chat_data + assert dp.chat_data[effective_new_chat_id]['key'] == "test" + def test_error_while_persisting(self, dp, caplog): class OwnPersistence(BasePersistence): def update(self, data): diff --git a/tests/test_persistence.py b/tests/test_persistence.py index 66ffbe7ff82..55ab9f1f15e 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -166,10 +166,10 @@ def update_callback_data(self, data): self.callback_data = data def drop_user_data(self, user_id): - pass + self.user_data.pop(user_id, None) def drop_chat_data(self, chat_id): - pass + self.chat_data.pop(chat_id, None) def update_conversation(self, name, key, new_state): raise NotImplementedError @@ -474,6 +474,15 @@ def delete_chat_data(chat_id): assert r.getMessage() == 'No error handlers are registered, logging exception.' assert r.levelname == 'ERROR' + def test_dispatcher_integration_migrate_chat_data(self, dp, bot_persistence): + dp.persistence = bot_persistence + dp.chat_data[1]['key'] = 'value' + dp.update_persistence() + assert bot_persistence.chat_data == {1: {'key': 'value'}} + + dp.migrate_chat_data(old_chat_id=1, new_chat_id=2) + assert bot_persistence.chat_data == {2: {'key': 'value'}} + @pytest.mark.parametrize( 'store_user_data', [True, False], ids=['store_user_data-True', 'store_user_data-False'] ) pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy