From 9c9b133ed2ac6e44d8ade8f2a77b7a5a75e969de Mon Sep 17 00:00:00 2001 From: DonalDuck004 <72563093+DonalDuck004@users.noreply.github.com> Date: Wed, 19 Jan 2022 21:33:04 +0100 Subject: [PATCH] Add Dispatcher.migrate_chat_data (#2848) Co-authored-by: Bibo-Joshi <22366557+Bibo-Joshi@users.noreply.github.com> Co-authored-by: Harshil <37377066+harshil21@users.noreply.github.com> --- telegram/ext/_dispatcher.py | 47 +++++++++++++++++++++++++++++++++++++ tests/test_dispatcher.py | 45 +++++++++++++++++++++++++++++++++++ tests/test_persistence.py | 13 ++++++++-- 3 files changed, 103 insertions(+), 2 deletions(-) diff --git a/telegram/ext/_dispatcher.py b/telegram/ext/_dispatcher.py index 3bba0df15..f635cf32c 100644 --- a/telegram/ext/_dispatcher.py +++ b/telegram/ext/_dispatcher.py @@ -55,6 +55,7 @@ from telegram.ext._utils.types import CCT, UD, CD, BD, BT, JQ, PT 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 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]): 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 6a077049a..eecc123b4 100644 --- a/tests/test_dispatcher.py +++ b/tests/test_dispatcher.py @@ -745,6 +745,51 @@ class TestDispatcher: 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 66ffbe7ff..55ab9f1f1 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -166,10 +166,10 @@ def bot_persistence(): 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 @@ class TestBasePersistence: 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'] )