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

@@ -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`.
 

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

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']
     )