Refactor persistence of Bot instances (#1994)

* Refactor persistence of bots

* Use BP.set_bot in Dispatcher

* Add documentation
This commit is contained in:
Bibo-Joshi 2020-07-13 21:52:26 +02:00
parent 19a4f9e53a
commit 2381724b7c
6 changed files with 261 additions and 4 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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