Refactor persistence of Bot instances (#1994)

* Refactor persistence of bots

* Use BP.set_bot in Dispatcher

* Add documentation

telegram/bot.py

@@ -3868,10 +3868,6 @@ class Bot(TelegramObject):
 
         return data
 
-    def __reduce__(self):
-        return (self.__class__, (self.token, self.base_url.replace(self.token, ''),
-                                 self.base_file_url.replace(self.token, '')))
-
     # camelCase aliases
     getMe = get_me
     """Alias for :attr:`get_me`"""

telegram/ext/basepersistence.py

@@ -19,6 +19,10 @@
 """This module contains the BasePersistence class."""
 
 from abc import ABC, abstractmethod
+from collections import defaultdict
+from copy import copy
+
+from telegram import Bot
 
 
 class BasePersistence(ABC):
@@ -37,6 +41,18 @@ class BasePersistence(ABC):
       must overwrite :meth:`get_conversations` and :meth:`update_conversation`.
     * :meth:`flush` will be called when the bot is shutdown.
 
+    Warning:
+        Persistence will try to replace :class:`telegram.Bot` instances by :attr:`REPLACED_BOT` and
+        insert the bot set with :meth:`set_bot` upon loading of the data. This is to ensure that
+        changes to the bot apply to the saved objects, too. If you change the bots token, this may
+        lead to e.g. ``Chat not found`` errors. For the limitations on replacing bots see
+        :meth:`replace_bot` and :meth:`insert_bot`.
+
+    Note:
+         :meth:`replace_bot` and :meth:`insert_bot` are used *independently* of the implementation
+         of the :meth:`update/get_*` methods, i.e. you don't need to worry about it while
+         implementing a custom persistence subclass.
+
     Attributes:
         store_user_data (:obj:`bool`): Optional, Whether user_data should be saved by this
             persistence class.
@@ -54,10 +70,128 @@ class BasePersistence(ABC):
             persistence class. Default is :obj:`True` .
     """
 
+    def __new__(cls, *args, **kwargs):
+        instance = super().__new__(cls)
+        get_user_data = instance.get_user_data
+        get_chat_data = instance.get_chat_data
+        get_bot_data = instance.get_bot_data
+        update_user_data = instance.update_user_data
+        update_chat_data = instance.update_chat_data
+        update_bot_data = instance.update_bot_data
+
+        def get_user_data_insert_bot():
+            return instance.insert_bot(get_user_data())
+
+        def get_chat_data_insert_bot():
+            return instance.insert_bot(get_chat_data())
+
+        def get_bot_data_insert_bot():
+            return instance.insert_bot(get_bot_data())
+
+        def update_user_data_replace_bot(user_id, data):
+            return update_user_data(user_id, instance.replace_bot(data))
+
+        def update_chat_data_replace_bot(chat_id, data):
+            return update_chat_data(chat_id, instance.replace_bot(data))
+
+        def update_bot_data_replace_bot(data):
+            return update_bot_data(instance.replace_bot(data))
+
+        instance.get_user_data = get_user_data_insert_bot
+        instance.get_chat_data = get_chat_data_insert_bot
+        instance.get_bot_data = get_bot_data_insert_bot
+        instance.update_user_data = update_user_data_replace_bot
+        instance.update_chat_data = update_chat_data_replace_bot
+        instance.update_bot_data = update_bot_data_replace_bot
+        return instance
+
     def __init__(self, store_user_data=True, store_chat_data=True, store_bot_data=True):
         self.store_user_data = store_user_data
         self.store_chat_data = store_chat_data
         self.store_bot_data = store_bot_data
+        self.bot = None
+
+    def set_bot(self, bot):
+        """Set the Bot to be used by this persistence instance.
+
+        Args:
+            bot (:class:`telegram.Bot`): The bot.
+        """
+        self.bot = bot
+
+    @classmethod
+    def replace_bot(cls, obj):
+        """
+        Replaces all instances of :class:`telegram.Bot` that occur within the passed object with
+        :attr:`REPLACED_BOT`. Currently, this handles objects of type ``list``, ``tuple``, ``set``,
+        ``frozenset``, ``dict``, ``defaultdict`` and objects that have a ``__dict__`` or
+        ``__slot__`` attribute.
+
+        Args:
+            obj (:obj:`object`): The object
+
+        Returns:
+            :obj:`obj`: Copy of the object with Bot instances replaced.
+        """
+        if isinstance(obj, Bot):
+            return cls.REPLACED_BOT
+        if isinstance(obj, (list, tuple, set, frozenset)):
+            return obj.__class__(cls.replace_bot(item) for item in obj)
+
+        new_obj = copy(obj)
+        if isinstance(obj, (dict, defaultdict)):
+            new_obj.clear()
+            for k, v in obj.items():
+                new_obj[cls.replace_bot(k)] = cls.replace_bot(v)
+            return new_obj
+        if hasattr(obj, '__dict__'):
+            for attr_name, attr in new_obj.__dict__.items():
+                setattr(new_obj, attr_name, cls.replace_bot(attr))
+            return new_obj
+        if hasattr(obj, '__slots__'):
+            for attr_name in new_obj.__slots__:
+                setattr(new_obj, attr_name,
+                        cls.replace_bot(cls.replace_bot(getattr(new_obj, attr_name))))
+            return new_obj
+
+        return obj
+
+    def insert_bot(self, obj):
+        """
+        Replaces all instances of :attr:`REPLACED_BOT` that occur within the passed object with
+        :attr:`bot`. Currently, this handles objects of type ``list``, ``tuple``, ``set``,
+        ``frozenset``, ``dict``, ``defaultdict`` and objects that have a ``__dict__`` or
+        ``__slot__`` attribute.
+
+        Args:
+            obj (:obj:`object`): The object
+
+        Returns:
+            :obj:`obj`: Copy of the object with Bot instances inserted.
+        """
+        if isinstance(obj, Bot):
+            return self.bot
+        if obj == self.REPLACED_BOT:
+            return self.bot
+        if isinstance(obj, (list, tuple, set, frozenset)):
+            return obj.__class__(self.insert_bot(item) for item in obj)
+
+        new_obj = copy(obj)
+        if isinstance(obj, (dict, defaultdict)):
+            new_obj.clear()
+            for k, v in obj.items():
+                new_obj[self.insert_bot(k)] = self.insert_bot(v)
+            return new_obj
+        if hasattr(obj, '__dict__'):
+            for attr_name, attr in new_obj.__dict__.items():
+                setattr(new_obj, attr_name, self.insert_bot(attr))
+            return new_obj
+        if hasattr(obj, '__slots__'):
+            for attr_name in obj.__slots__:
+                setattr(new_obj, attr_name,
+                        self.insert_bot(self.insert_bot(getattr(new_obj, attr_name))))
+            return new_obj
+        return obj
 
     @abstractmethod
     def get_user_data(self):
@@ -149,3 +283,6 @@ class BasePersistence(ABC):
         is not of any importance just pass will be sufficient.
         """
         pass
+
+    REPLACED_BOT = 'bot_instance_replaced_by_ptb_persistence'
+    """:obj:`str`: Placeholder for :class:`telegram.Bot` instances replaced in saved data."""

telegram/ext/dictpersistence.py

@@ -39,6 +39,15 @@ class DictPersistence(BasePersistence):
         because ``DictPersistence`` is mainly intended as starting point for custom persistence
         classes that need to JSON-serialize the stored data before writing them to file/database.
 
+    Warning:
+        :class:`DictPersistence` will try to replace :class:`telegram.Bot` instances by
+        :attr:`REPLACED_BOT` and insert the bot set with
+        :meth:`telegram.ext.BasePersistence.set_bot` upon loading of the data. This is to ensure
+        that changes to the bot apply to the saved objects, too. If you change the bots token, this
+        may lead to e.g. ``Chat not found`` errors. For the limitations on replacing bots see
+        :meth:`telegram.ext.BasePersistence.replace_bot` and
+        :meth:`telegram.ext.BasePersistence.insert_bot`.
+
     Attributes:
         store_user_data (:obj:`bool`): Whether user_data should be saved by this
             persistence class.

telegram/ext/dispatcher.py

@@ -155,6 +155,7 @@ class Dispatcher:
             if not isinstance(persistence, BasePersistence):
                 raise TypeError("persistence must be based on telegram.ext.BasePersistence")
             self.persistence = persistence
+            self.persistence.set_bot(self.bot)
             if self.persistence.store_user_data:
                 self.user_data = self.persistence.get_user_data()
                 if not isinstance(self.user_data, defaultdict):

telegram/ext/picklepersistence.py

@@ -27,6 +27,15 @@ from telegram.ext import BasePersistence
 class PicklePersistence(BasePersistence):
     """Using python's builtin pickle for making you bot persistent.
 
+    Warning:
+        :class:`PicklePersistence` will try to replace :class:`telegram.Bot` instances by
+        :attr:`REPLACED_BOT` and insert the bot set with
+        :meth:`telegram.ext.BasePersistence.set_bot` upon loading of the data. This is to ensure
+        that changes to the bot apply to the saved objects, too. If you change the bots token, this
+        may lead to e.g. ``Chat not found`` errors. For the limitations on replacing bots see
+        :meth:`telegram.ext.BasePersistence.replace_bot` and
+        :meth:`telegram.ext.BasePersistence.insert_bot`.
+
     Attributes:
         filename (:obj:`str`): The filename for storing the pickle files. When :attr:`single_file`
             is :obj:`False` this will be used as a prefix.

tests/test_persistence.py

@@ -395,6 +395,111 @@ class TestBasePersistence:
             dp.process_update(MyUpdate())
         assert 'An uncaught error was raised while processing the update' not in caplog.text
 
+    def test_bot_replace_insert_bot(self, bot):
+
+        class BotPersistence(BasePersistence):
+            def __init__(self):
+                super().__init__()
+                self.bot_data = None
+                self.chat_data = defaultdict(dict)
+                self.user_data = defaultdict(dict)
+
+            def get_bot_data(self):
+                return self.bot_data
+
+            def get_chat_data(self):
+                return self.chat_data
+
+            def get_user_data(self):
+                return self.user_data
+
+            def get_conversations(self, name):
+                raise NotImplementedError
+
+            def update_bot_data(self, data):
+                self.bot_data = data
+
+            def update_chat_data(self, chat_id, data):
+                self.chat_data[chat_id] = data
+
+            def update_user_data(self, user_id, data):
+                self.user_data[user_id] = data
+
+            def update_conversation(self, name, key, new_state):
+                raise NotImplementedError
+
+        class CustomSlottedClass:
+            __slots__ = ('bot',)
+
+            def __init__(self):
+                self.bot = bot
+
+            def __eq__(self, other):
+                if isinstance(other, CustomSlottedClass):
+                    return self.bot is other.bot
+                return False
+
+        class CustomClass:
+            def __init__(self):
+                self.bot = bot
+                self.slotted_object = CustomSlottedClass()
+                self.list_ = [1, 2, bot]
+                self.tuple_ = tuple(self.list_)
+                self.set_ = set(self.list_)
+                self.frozenset_ = frozenset(self.list_)
+                self.dict_ = {item: item for item in self.list_}
+                self.defaultdict_ = defaultdict(dict, self.dict_)
+
+            @staticmethod
+            def replace_bot():
+                cc = CustomClass()
+                cc.bot = BasePersistence.REPLACED_BOT
+                cc.slotted_object.bot = BasePersistence.REPLACED_BOT
+                cc.list_ = [1, 2, BasePersistence.REPLACED_BOT]
+                cc.tuple_ = tuple(cc.list_)
+                cc.set_ = set(cc.list_)
+                cc.frozenset_ = frozenset(cc.list_)
+                cc.dict_ = {item: item for item in cc.list_}
+                cc.defaultdict_ = defaultdict(dict, cc.dict_)
+                return cc
+
+            def __eq__(self, other):
+                if isinstance(other, CustomClass):
+                    # print(self.__dict__)
+                    # print(other.__dict__)
+                    return (self.bot == other.bot
+                            and self.slotted_object == other.slotted_object
+                            and self.list_ == other.list_
+                            and self.tuple_ == other.tuple_
+                            and self.set_ == other.set_
+                            and self.frozenset_ == other.frozenset_
+                            and self.dict_ == other.dict_
+                            and self.defaultdict_ == other.defaultdict_)
+                return False
+
+        persistence = BotPersistence()
+        persistence.set_bot(bot)
+        cc = CustomClass()
+
+        persistence.update_bot_data({1: cc})
+        assert persistence.bot_data[1].bot == BasePersistence.REPLACED_BOT
+        assert persistence.bot_data[1] == cc.replace_bot()
+
+        persistence.update_chat_data(123, {1: cc})
+        assert persistence.chat_data[123][1].bot == BasePersistence.REPLACED_BOT
+        assert persistence.chat_data[123][1] == cc.replace_bot()
+
+        persistence.update_user_data(123, {1: cc})
+        assert persistence.user_data[123][1].bot == BasePersistence.REPLACED_BOT
+        assert persistence.user_data[123][1] == cc.replace_bot()
+
+        assert persistence.get_bot_data()[1] == cc
+        assert persistence.get_bot_data()[1].bot is bot
+        assert persistence.get_chat_data()[123][1] == cc
+        assert persistence.get_chat_data()[123][1].bot is bot
+        assert persistence.get_user_data()[123][1] == cc
+        assert persistence.get_user_data()[123][1].bot is bot
+
 
 @pytest.fixture(scope='function')
 def pickle_persistence():