mirror of
https://github.com/python-telegram-bot/python-telegram-bot.git
synced 2025-03-17 04:39:55 +01:00
Refactor persistence of Bot instances (#1994)
* Refactor persistence of bots * Use BP.set_bot in Dispatcher * Add documentation
This commit is contained in:
parent
19a4f9e53a
commit
2381724b7c
6 changed files with 261 additions and 4 deletions
|
@ -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`"""
|
||||
|
|
|
@ -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."""
|
||||
|
|
|
@ -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.
|
||||
|
|
|
@ -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):
|
||||
|
|
|
@ -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.
|
||||
|
|
|
@ -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():
|
||||
|
|
Loading…
Add table
Reference in a new issue