Skip to content

Commit f6b663f

Browse files
Bibo-JoshiEldinnie
Bibo-Joshi
and
Eldinnie
authored
bot_data as global memory (V12 version of python-telegram-bot#1322) (python-telegram-bot#1325)
* Update AUTHORS.rst * Update AUTHORS.rst * Add bot_data to CallbackContext as global memory * Minor fixes in docstrings * Incorp. req. changes, Flake8 Fixes * Persist before stop * Fix CI errors * Implement python-telegram-bot#1342 for bot_data * Add check pickle_persistence_only_bot similar to python-telegram-bot#1462 * Fix test_persistence * Try dispatching error before logging it * Fix test Co-authored-by: Eldinnie <[email protected]>
1 parent 43bfebb commit f6b663f

21 files changed

Lines changed: 427 additions & 33 deletions

‎telegram/ext/basepersistence.py‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ class BasePersistence(object):
2525
2626
All relevant methods must be overwritten. This means:
2727
28+
* If :attr:`store_bot_data` is ``True`` you must overwrite :meth:`get_bot_data` and
29+
:meth:`update_bot_data`.
2830
* If :attr:`store_chat_data` is ``True`` you must overwrite :meth:`get_chat_data` and
2931
:meth:`update_chat_data`.
3032
* If :attr:`store_user_data` is ``True`` you must overwrite :meth:`get_user_data` and
@@ -38,17 +40,22 @@ class BasePersistence(object):
3840
persistence class.
3941
store_chat_data (:obj:`bool`): Optional. Whether chat_data should be saved by this
4042
persistence class.
43+
store_bot_data (:obj:`bool`): Optional. Whether bot_data should be saved by this
44+
persistence class.
4145
4246
Args:
4347
store_user_data (:obj:`bool`, optional): Whether user_data should be saved by this
4448
persistence class. Default is ``True``.
4549
store_chat_data (:obj:`bool`, optional): Whether chat_data should be saved by this
4650
persistence class. Default is ``True`` .
51+
store_bot_data (:obj:`bool`, optional): Whether bot_data should be saved by this
52+
persistence class. Default is ``True`` .
4753
"""
4854

49-
def __init__(self, store_user_data=True, store_chat_data=True):
55+
def __init__(self, store_user_data=True, store_chat_data=True, store_bot_data=True):
5056
self.store_user_data = store_user_data
5157
self.store_chat_data = store_chat_data
58+
self.store_bot_data = store_bot_data
5259

5360
def get_user_data(self):
5461
""""Will be called by :class:`telegram.ext.Dispatcher` upon creation with a
@@ -70,6 +77,16 @@ def get_chat_data(self):
7077
"""
7178
raise NotImplementedError
7279

80+
def get_bot_data(self):
81+
""""Will be called by :class:`telegram.ext.Dispatcher` upon creation with a
82+
persistence object. It should return the bot_data if stored, or an empty
83+
``dict``.
84+
85+
Returns:
86+
:obj:`defaultdict`: The restored bot data.
87+
"""
88+
raise NotImplementedError
89+
7390
def get_conversations(self, name):
7491
""""Will be called by :class:`telegram.ext.Dispatcher` when a
7592
:class:`telegram.ext.ConversationHandler` is added if
@@ -111,7 +128,16 @@ def update_chat_data(self, chat_id, data):
111128
112129
Args:
113130
chat_id (:obj:`int`): The chat the data might have been changed for.
114-
data (:obj:`dict`): The :attr:`telegram.ext.dispatcher.chat_data` [user_id].
131+
data (:obj:`dict`): The :attr:`telegram.ext.dispatcher.chat_data` [chat_id].
132+
"""
133+
raise NotImplementedError
134+
135+
def update_bot_data(self, data):
136+
"""Will be called by the :class:`telegram.ext.Dispatcher` after a handler has
137+
handled an update.
138+
139+
Args:
140+
data (:obj:`dict`): The :attr:`telegram.ext.dispatcher.bot_data` .
115141
"""
116142
raise NotImplementedError
117143

‎telegram/ext/callbackcontext.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ class CallbackContext(object):
4343
that you think you added will not be present.
4444
4545
Attributes:
46+
bot_data (:obj:`dict`, optional): A dict that can be used to keep any data in. For each
47+
update it will be the same ``dict``.
4648
chat_data (:obj:`dict`, optional): A dict that can be used to keep any data in. For each
4749
update from the same chat id it will be the same ``dict``.
4850
@@ -80,6 +82,7 @@ def __init__(self, dispatcher):
8082
raise ValueError('CallbackContext should not be used with a non context aware '
8183
'dispatcher!')
8284
self._dispatcher = dispatcher
85+
self._bot_data = dispatcher.bot_data
8386
self._chat_data = None
8487
self._user_data = None
8588
self.args = None
@@ -92,6 +95,15 @@ def dispatcher(self):
9295
""":class:`telegram.ext.Dispatcher`: The dispatcher associated with this context."""
9396
return self._dispatcher
9497

98+
@property
99+
def bot_data(self):
100+
return self._bot_data
101+
102+
@bot_data.setter
103+
def bot_data(self, value):
104+
raise AttributeError("You can not assign a new value to bot_data, see "
105+
"https://git.io/fjxKe")
106+
95107
@property
96108
def chat_data(self):
97109
return self._chat_data
@@ -119,6 +131,7 @@ def from_error(cls, update, error, dispatcher):
119131
@classmethod
120132
def from_update(cls, update, dispatcher):
121133
self = cls(dispatcher)
134+
122135
if update is not None and isinstance(update, Update):
123136
chat = update.effective_chat
124137
user = update.effective_user

‎telegram/ext/dictpersistence.py‎

Lines changed: 64 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,36 +31,51 @@
3131

3232

3333
class DictPersistence(BasePersistence):
34-
"""Using python's dicts and json for making you bot persistent.
34+
"""Using python's dicts and json for making your bot persistent.
3535
3636
Attributes:
3737
store_user_data (:obj:`bool`): Whether user_data should be saved by this
3838
persistence class.
3939
store_chat_data (:obj:`bool`): Whether chat_data should be saved by this
4040
persistence class.
41+
store_bot_data (:obj:`bool`): Whether bot_data should be saved by this
42+
persistence class.
4143
4244
Args:
4345
store_user_data (:obj:`bool`, optional): Whether user_data should be saved by this
4446
persistence class. Default is ``True``.
4547
store_chat_data (:obj:`bool`, optional): Whether user_data should be saved by this
4648
persistence class. Default is ``True``.
49+
store_bot_data (:obj:`bool`, optional): Whether bot_data should be saved by this
50+
persistence class. Default is ``True`` .
4751
user_data_json (:obj:`str`, optional): Json string that will be used to reconstruct
4852
user_data on creating this persistence. Default is ``""``.
4953
chat_data_json (:obj:`str`, optional): Json string that will be used to reconstruct
5054
chat_data on creating this persistence. Default is ``""``.
55+
bot_data_json (:obj:`str`, optional): Json string that will be used to reconstruct
56+
bot_data on creating this persistence. Default is ``""``.
5157
conversations_json (:obj:`str`, optional): Json string that will be used to reconstruct
5258
conversation on creating this persistence. Default is ``""``.
5359
"""
5460

55-
def __init__(self, store_user_data=True, store_chat_data=True, user_data_json='',
56-
chat_data_json='', conversations_json=''):
57-
self.store_user_data = store_user_data
58-
self.store_chat_data = store_chat_data
61+
def __init__(self,
62+
store_user_data=True,
63+
store_chat_data=True,
64+
store_bot_data=True,
65+
user_data_json='',
66+
chat_data_json='',
67+
bot_data_json='',
68+
conversations_json=''):
69+
super(DictPersistence, self).__init__(store_user_data=store_user_data,
70+
store_chat_data=store_chat_data,
71+
store_bot_data=store_bot_data)
5972
self._user_data = None
6073
self._chat_data = None
74+
self._bot_data = None
6175
self._conversations = None
6276
self._user_data_json = None
6377
self._chat_data_json = None
78+
self._bot_data_json = None
6479
self._conversations_json = None
6580
if user_data_json:
6681
try:
@@ -74,6 +89,14 @@ def __init__(self, store_user_data=True, store_chat_data=True, user_data_json=''
7489
self._chat_data_json = chat_data_json
7590
except (ValueError, AttributeError):
7691
raise TypeError("Unable to deserialize chat_data_json. Not valid JSON")
92+
if bot_data_json:
93+
try:
94+
self._bot_data = json.loads(bot_data_json)
95+
self._bot_data_json = bot_data_json
96+
except (ValueError, AttributeError):
97+
raise TypeError("Unable to deserialize bot_data_json. Not valid JSON")
98+
if not isinstance(self._bot_data, dict):
99+
raise TypeError("bot_data_json must be serialized dict")
77100

78101
if conversations_json:
79102
try:
@@ -108,6 +131,19 @@ def chat_data_json(self):
108131
else:
109132
return json.dumps(self.chat_data)
110133

134+
@property
135+
def bot_data(self):
136+
""":obj:`dict`: The bot_data as a dict"""
137+
return self._bot_data
138+
139+
@property
140+
def bot_data_json(self):
141+
""":obj:`str`: The bot_data serialized as a JSON-string."""
142+
if self._bot_data_json:
143+
return self._bot_data_json
144+
else:
145+
return json.dumps(self.bot_data)
146+
111147
@property
112148
def conversations(self):
113149
""":obj:`dict`: The conversations as a dict"""
@@ -145,6 +181,18 @@ def get_chat_data(self):
145181
self._chat_data = defaultdict(dict)
146182
return deepcopy(self.chat_data)
147183

184+
def get_bot_data(self):
185+
"""Returns the bot_data created from the ``bot_data_json`` or an empty dict.
186+
187+
Returns:
188+
:obj:`defaultdict`: The restored user data.
189+
"""
190+
if self.bot_data:
191+
pass
192+
else:
193+
self._bot_data = {}
194+
return deepcopy(self.bot_data)
195+
148196
def get_conversations(self, name):
149197
"""Returns the conversations created from the ``conversations_json`` or an empty
150198
defaultdict.
@@ -194,3 +242,14 @@ def update_chat_data(self, chat_id, data):
194242
return
195243
self._chat_data[chat_id] = data
196244
self._chat_data_json = None
245+
246+
def update_bot_data(self, data):
247+
"""Will update the bot_data (if changed).
248+
249+
Args:
250+
data (:obj:`dict`): The :attr:`telegram.ext.dispatcher.bot_data`.
251+
"""
252+
if self._bot_data == data:
253+
return
254+
self._bot_data = data.copy()
255+
self._bot_data_json = None

‎telegram/ext/dispatcher.py‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ class Dispatcher(object):
7979
decorator.
8080
user_data (:obj:`defaultdict`): A dictionary handlers can use to store data for the user.
8181
chat_data (:obj:`defaultdict`): A dictionary handlers can use to store data for the chat.
82+
bot_data (:obj:`dict`): A dictionary handlers can use to store data for the bot.
8283
persistence (:class:`telegram.ext.BasePersistence`): Optional. The persistence class to
8384
store data that should be persistent over restarts
8485
@@ -121,8 +122,8 @@ def __init__(self,
121122
TelegramDeprecationWarning, stacklevel=3)
122123

123124
self.user_data = defaultdict(dict)
124-
""":obj:`dict`: A dictionary handlers can use to store data for the user."""
125125
self.chat_data = defaultdict(dict)
126+
self.bot_data = {}
126127
if persistence:
127128
if not isinstance(persistence, BasePersistence):
128129
raise TypeError("persistence should be based on telegram.ext.BasePersistence")
@@ -135,6 +136,10 @@ def __init__(self,
135136
self.chat_data = self.persistence.get_chat_data()
136137
if not isinstance(self.chat_data, defaultdict):
137138
raise ValueError("chat_data must be of type defaultdict")
139+
if self.persistence.store_bot_data:
140+
self.bot_data = self.persistence.get_bot_data()
141+
if not isinstance(self.bot_data, dict):
142+
raise ValueError("bot_data must be of type dict")
138143
else:
139144
self.persistence = None
140145

@@ -327,6 +332,17 @@ def persist_update(update):
327332
328333
"""
329334
if self.persistence and isinstance(update, Update):
335+
if self.persistence.store_bot_data:
336+
try:
337+
self.persistence.update_bot_data(self.bot_data)
338+
except Exception as e:
339+
try:
340+
self.dispatch_error(update, e)
341+
except Exception:
342+
message = 'Saving bot data raised an error and an ' \
343+
'uncaught error was raised while handling ' \
344+
'the error with an error_handler'
345+
self.logger.exception(message)
330346
if self.persistence.store_chat_data and update.effective_chat:
331347
chat_id = update.effective_chat.id
332348
try:
@@ -456,9 +472,11 @@ def remove_handler(self, handler, group=DEFAULT_GROUP):
456472
self.groups.remove(group)
457473

458474
def update_persistence(self):
459-
"""Update :attr:`user_data` and :attr:`chat_data` in :attr:`persistence`.
475+
"""Update :attr:`user_data`, :attr:`chat_data` and :attr:`bot_data` in :attr:`persistence`.
460476
"""
461477
if self.persistence:
478+
if self.persistence.store_bot_data:
479+
self.persistence.update_bot_data(self.bot_data)
462480
if self.persistence.store_chat_data:
463481
for chat_id in self.chat_data:
464482
self.persistence.update_chat_data(chat_id, self.chat_data[chat_id])

0 commit comments

Comments
 (0)