diff --git a/telegram/bot.py b/telegram/bot.py index aa863f79a56..e38cafe0cdb 100644 --- a/telegram/bot.py +++ b/telegram/bot.py @@ -3811,10 +3811,6 @@ def to_dict(self): 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`""" diff --git a/telegram/ext/basepersistence.py b/telegram/ext/basepersistence.py index b4004a7c33f..4e507d61e2c 100644 --- a/telegram/ext/basepersistence.py +++ b/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 ``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 @@ def flush(self): 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.""" diff --git a/telegram/ext/dictpersistence.py b/telegram/ext/dictpersistence.py index 3d18aa14883..72323928f21 100644 --- a/telegram/ext/dictpersistence.py +++ b/telegram/ext/dictpersistence.py @@ -33,6 +33,15 @@ class DictPersistence(BasePersistence): """Using python's dicts and json for making your bot persistent. + 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. diff --git a/telegram/ext/dispatcher.py b/telegram/ext/dispatcher.py index 5dd61ed28bc..5c4cdaaf490 100644 --- a/telegram/ext/dispatcher.py +++ b/telegram/ext/dispatcher.py @@ -125,6 +125,7 @@ def __init__(self, if not isinstance(persistence, BasePersistence): raise TypeError("persistence should 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): diff --git a/telegram/ext/picklepersistence.py b/telegram/ext/picklepersistence.py index 2c484a7db36..24091145647 100644 --- a/telegram/ext/picklepersistence.py +++ b/telegram/ext/picklepersistence.py @@ -27,6 +27,15 @@ 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 false this will be used as a prefix. diff --git a/tests/test_persistence.py b/tests/test_persistence.py index 9e7178d07fb..fec89d06afd 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -295,6 +295,111 @@ class MyUpdate: 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():