8000 Refactor persistence of Bot instances (#1994) · Konano/python-telegram-bot@2381724 · GitHub
[go: up one dir, main page]
Skip to content

Commit 2381724

Browse files
Bibo-JoshiBibo-Joshi
committed
Refactor persistence of Bot instances (python-telegram-bot#1994)
* Refactor persistence of bots * Use BP.set_bot in Dispatcher * Add documentation
1 parent 19a4f9e commit 2381724

File tree

6 files changed

+261
-4
lines changed

6 files changed

+261
-4
lines changed

‎telegram/bot.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3868,10 +3868,6 @@ def to_dict(self):
38683868

38693869
return data
38703870

3871-
def __reduce__(self):
3872-
return (self.__class__, (self.token, self.base_url.replace(self.token, ''),
3873-
self.base_file_url.replace(self.token, '')))
3874-
38753871
# camelCase aliases
38763872
getMe = get_me
38773873
"""Alias for :attr:`get_me`"""

‎telegram/ext/basepersistence.py

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,10 @@
1919
"""This module contains the BasePersistence class."""
2020

2121
from abc import ABC, abstractmethod
22+
from collections import defaultdict
23+
from copy import copy
24+
25+
from telegram import Bot
2226

2327

2428
class BasePersistence(ABC):
@@ -37,6 +41,18 @@ class BasePersistence(ABC):
3741
must overwrite :meth:`get_conversations` and :meth:`update_conversation`.
3842
* :meth:`flush` will be called when the bot is shutdown.
3943
44+
Warning:
45+
Persistence will try to replace :class:`telegram.Bot` instances by :attr:`REPLACED_BOT` and
46+
insert the bot set with :meth:`set_bot` upon loading of the data. This is to ensure that
47+
changes to the bot apply to the saved objects, too. If you change the bots token, this may
48+
lead to e.g. ``Chat not found`` errors. For the limitations on replacing bots see
49+
:meth:`replace_bot` and :meth:`insert_bot`.
50+
51+
Note:
52+
:meth:`replace_bot` and :meth:`insert_bot` are used *independently* of the implementation
53+
of the :meth:`update/get_*` methods, i.e. you don't need to worry about it while
54+
implementing a custom persistence subclass.
55+
4056
Attributes:
4157
store_user_data (:obj:`bool`): Optional, Whether user_data should be saved by this
4258
persistence class.
@@ -54,10 +70,128 @@ class BasePersistence(ABC):
5470
persistence class. Default is :obj:`True` .
5571
"""
5672

73+
def __new__(cls, *args, **kwargs):
74+
instance = super().__new__(cls)
75+
get_user_data = instance.get_user_data
76+
get_chat_data = instance.get_chat_data
77+
get_bot_data = instance.get_bot_data
78+
update_user_data = instance.update_user_data
79+
update_chat_data = instance.update_chat_data
80+
update_bot_data = instance.update_bot_data
81+
82+
def get_user_data_insert_bot():
83+
return instance.insert_bot(get_user_data())
84+
85+
def get_chat_data_insert_bot():
86+
return instance.insert_bot(get_chat_data())
87+
88+
def get_bot_data_insert_bot():
89+
return instance.insert_bot(get_bot_data())
90+
91+
def update_user_data_replace_bot(user_id, data):
92+
return update_user_data(user_id, instance.replace_bot(data))
93+
94+
def update_chat_data_replace_bot(chat_id, data):
95+
return update_chat_data(chat_id, instance.replace_bot(data))
96+
97+
def update_bot_data_replace_bot(data):
98+
return update_bot_data(instance.replace_bot(data))
99+
100+
instance.get_user_data = get_user_data_insert_bot
101+
instance.get_chat_data = get_chat_data_insert_bot
102+
instance.get_bot_data = get_bot_data_insert_bot
103+
instance.update_user_data = update_user_data_replace_bot
104+
instance.update_chat_data = update_chat_data_replace_bot
105+
instance.update_bot_data = update_bot_data_replace_bot
106+
return instance
107+
57108
def __init__(self, store_user_data=True, store_chat_data=True, store_bot_data=True):
58109
self.store_user_data = store_user_data
59110
self.store_chat_data = store_chat_data
60111
self.store_bot_data = store_bot_data
112+
self.bot = None
113+
114+
def set_bot(self, bot):
115+
"""Set the Bot to be used by this persistence instance.
116+
117+
Args:
118+
bot (:class:`telegram.Bot`): The bot.
119+
"""
120+
self.bot = bot
121+
122+
@classmethod
123+
def replace_bot(cls, obj):
124+
"""
125+
Replaces all instances of :class:`telegram.Bot` that occur within the passed object with
126+
:attr:`REPLACED_BOT`. Currently, this handles objects of type ``list``, ``tuple``, ``set``,
127+
``frozenset``, ``dict``, ``defaultdict`` and objects that have a ``__dict__`` or
128+
``__slot__`` attribute.
129+
130+
Args:
131+
obj (:obj:`object`): The object
132+
133+
Returns:
134+
:obj:`obj`: Copy of the object with Bot instances replaced.
135+
"""
136+
if isinstance(obj, Bot):
137+
return cls.REPLACED_BOT
138+
if isinstance(obj, (list, tuple, set, frozenset)):
139+
return obj.__class__(cls.replace_bot(item) for item in obj)
140+
141+
new_obj = copy(obj)
142+
if isinstance(obj, (dict, defaultdict)):
143+
new_obj.clear()
144+
for k, v in obj.items():
145+
new_obj[cls.replace_bot(k)] = cls.replace_bot(v)
146+
return new_obj
147+
if hasattr(obj, '__dict__'):
148+
for attr_name, attr in new_obj.__dict__.items():
149+
setattr(new_obj, attr_name, cls.replace_bot(attr))
150+
return new_obj
151+
if hasattr(obj, '__slots__'):
152+
for attr_name in new_obj.__slots__:
153+
setattr(new_obj, attr_name,
154+
cls.replace_bot(cls.replace_bot(getattr(new_obj, attr_name))))
155+
return new_obj
156+
157+
return obj
158+
159+
def insert_bot(self, obj):
160+
"""
161+
Replaces all instances of :attr:`REPLACED_BOT` that occur within the passed object with
162+
:attr:`bot`. Currently, this handles objects of type ``list``, ``tuple``, ``set``,
163+
``frozenset``, ``dict``, ``defaultdict`` and objects that have a ``__dict__`` or
164+
``__slot__`` attribute.
165+
166+
Args:
167+
obj (:obj:`object`): The object
168+
169+
Returns:
170+
:obj:`obj`: Copy of the object with Bot instances inserted.
171+
"""
172+
if isinstance(obj, Bot):
173+
return self.bot
174+
if obj == self.REPLACED_BOT:
175+
return self.bot
176+
if isinstance(obj, (list, tuple, set, frozenset)):
177+
return obj.__class__(self.insert_bot(item) for item in obj)
178+
179+
new_obj = copy(obj)
180+
if isinstance(obj, (dict, defaultdict)):
181+
new_obj.clear()
182+
for k, v in obj.items():
183+
new_obj[self.insert_bot(k)] = self.insert_bot(v)
184+
return new_obj
185+
if hasattr(obj, '__dict__'):
186+
for attr_name, attr in new_obj.__dict__.items():
187+
setattr(new_obj, attr_name, self.insert_bot(attr))
188+
return new_obj
189+
if hasattr(obj, '__slots__'):
190+
for attr_name in obj.__slots__:
191+
setattr(new_obj, attr_name,
192+
self.insert_bot(self.insert_bot(getattr(new_obj, attr_name))))
193+
return new_obj
194+
return obj
61195

62196
@abstractmethod
63197
def get_user_data(self):
@@ -149,3 +283,6 @@ def flush(self):
149283
is not of any importance just pass will be sufficient.
150284
"""
151285
pass
286+
287+
REPLACED_BOT = 'bot_instance_replaced_by_ptb_persistence'
288+
""":obj:`str`: Placeholder for :class:`telegram.Bot` instances replaced in saved data."""

‎telegram/ext/dictpersistence.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,15 @@ class DictPersistence(BasePersistence):
3939
because ``DictPersistence`` is mainly intended as starting point for custom persistence
4040
classes that need to JSON-serialize the stored data before writing them to file/database.
4141
42+
Warning:
43+
:class:`DictPersistence` will try to replace :class:`telegram.Bot` instances by
44+
:attr:`REPLACED_BOT` and insert the bot set with
45+
:meth:`telegram.ext.BasePersistence.set_bot` upon loading of the data. This is to ensure
46+
that changes to the bot apply to the saved objects, too. If you change the bots token, this
47+
may lead to e.g. ``Chat not found`` errors. For the limitations on replacing bots see
48+
:meth:`telegram.ext.BasePersistence.replace_bot` and
49+
:meth:`telegram.ext.BasePersistence.insert_bot`.
50+
4251
Attributes:
4352
store_user_data (:obj:`bool`): Whether user_data should be saved by this
4453
persistence class.

‎telegram/ext/dispatcher.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,7 @@ def __init__(self,
155155
if not isinstance(persistence, BasePersistence):
156156
raise TypeError("persistence must be based on telegram.ext.BasePersistence")
157157
self.persistence = persistence
158+
self.persistence.set_bot(self.bot)
158159
if self.persistence.store_user_data:
159160
self.user_data = self.persistence.get_user_data()
160161
if not isinstance(self.user_data, defaultdict):

‎telegram/ext/picklepersistence.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,15 @@
2727
class PicklePersistence(BasePersistence):
2828
"""Using python's builtin pickle for making you bot persistent.
2929
30+
Warning:
31+
:class:`PicklePersistence` will try to replace :class:`telegram.Bot` instances by
32+
:attr:`REPLACED_BOT` and insert the bot set with
33+
:meth:`telegram.ext.BasePersistence.set_bot` upon loading of the data. This is to ensure
34+
that changes to the bot apply to the saved objects, too. If you change the bots token, this
35+
may lead to e.g. ``Chat not found`` errors. For the limitations on replacing bots see
36+
:meth:`telegram.ext.BasePersistence.replace_bot` and
37+
:meth:`telegram.ext.BasePersistence.insert_bot`.
38+
3039
Attributes:
3140
filename (:obj:`str`): The filename for storing the pickle files. When :attr:`single_file`
3241
is :obj:`False` this will be used as a prefix.

‎tests/test_persistence.py

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -395,6 +395,111 @@ class MyUpdate:
395395
dp.process_update(MyUpdate())
396396
assert 'An uncaught error was raised while processing the update' not in caplog.text
397397

398+
def test_bot_replace_insert_bot(self, bot):
399+
400+
class BotPersistence(BasePersistence):
401+
def __init__(self):
402+
super().__init__()
403+
self.bot_data = None
404+
self.chat_data = defaultdict(dict)
405+
self.user_data = defaultdict(dict)
406+
407+
def get_bot_data(self):
408+
return self.bot_data
409+
410+
def get_chat_data(self):
411+
return self.chat_data
412+
413+
def get_user_data(self):
414+
return self.user_data
415+
416+
def get_conversations(self, name):
417+
raise NotImplementedError
418+
419+
def update_bot_data(self, data):
420+
self.bot_data = data
421+
422+
def update_chat_data(self, chat_id, data):
423+
self.chat_data[chat_id] = data
424+
425+
def update_user_data(self, user_id, data):
426+
self.user_data[user_id] = data
427+
428+
def update_conversation(self, name, key, new_state):
429+
raise NotImplementedError
430+
431+
class CustomSlottedClass:
432+
__slots__ = ('bot',)
433+
434+
def __init__(self):
435+
self.bot = bot
436+
437+
def __eq__(self, other):
438+
if isinstance(other, CustomSlottedClass):
439+
return self.bot is other.bot
440+
return False
441+
442+
class CustomClass:
443+
def __init__(self):
444+
self.bot = bot
445+
self.slotted_object = CustomSlottedClass()
446+
self.list_ = [1, 2, bot]
447+
self.tuple_ = tuple(self.list_)
448+
self.set_ = set(self.list_)
449+
self.frozenset_ = frozenset(self.list_)
450+
self.dict_ = {item: item for item in self.list_}
451+
self.defaultdict_ = defaultdict(dict, self.dict_)
452+
453+
@staticmethod
454+
def replace_bot():
455+
cc = CustomClass()
456+
cc.bot = BasePersistence.REPLACED_BOT
457+
cc.slotted_object.bot = BasePersistence.REPLACED_BOT
458+
cc.list_ = [1, 2, BasePersistence.REPLACED_BOT]
459+
cc.tuple_ = tuple(cc.list_)
460+
cc.set_ = set(cc.list_)
461+
cc.frozenset_ = frozenset(cc.list_)
462+
cc.dict_ = {item: item for item in cc.list_}
463+
cc.defaultdict_ = defaultdict(dict, cc.dict_)
464+
return cc
465+
466+
def __eq__(self, other):
467+
if isinstance(other, CustomClass):
468+
# print(self.__dict__)
469+
# print(other.__dict__)
470+
return (self.bot == other.bot
471+
and self.slotted_object == other.slotted_object
472+
and self.list_ == other.list_
473+
and self.tuple_ == other.tuple_
474+
and self.set_ == other.set_
475+
and self.frozenset_ == other.frozenset_
476+
and self.dict_ == other.dict_
477+
and self.defaultdict_ == other.defaultdict_)
478+
return False
479+
480+
persistence = BotPersistence()
481+
persistence.set_bot(bot)
482+
cc = CustomClass()
483+
484+
persistence.update_bot_data({1: cc})
485+
assert persistence.bot_data[1].bot == BasePersistence.REPLACED_BOT
486+
assert persistence.bot_data[1] == cc.replace_bot()
487+
488+
persistence.update_chat_data(123, {1: cc})
489+
assert persistence.chat_data[123][1].bot == BasePersistence.REPLACED_BOT
490+
assert persistence.chat_data[123][1] == cc.replace_bot()
491+
492+
persistence.update_user_data(123, {1: cc})
493+
assert persistence.user_data[123][1].bot == BasePersistence.REPLACED_BOT
494+
assert persistence.user_data[123][1] == cc.replace_bot()
495+
496+
assert persistence.get_bot_data()[1] == cc
497+
assert persistence.get_bot_data()[1].bot is bot
498+
assert persistence.get_chat_data()[123][1] == cc
499+
assert persistence.get_chat_data()[123][1].bot is bot
500+
assert persistence.get_user_data()[123][1] == cc
501+
assert persistence.get_user_data()[123][1].bot is bot
502+
398503

399504
@pytest.fixture(scope='function')
400505
def pickle_persistence():

0 commit comments

Comments
 (0)
0