Skip to content

Commit 4f19acf

Browse files
EldinnieEldinnie
committed
PickelPersistence and start tests
1 parent ecf88fe commit 4f19acf

File tree

4 files changed

+553
-11
lines changed

4 files changed

+553
-11
lines changed

‎telegram/ext/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
"""Extensions over the Telegram Bot API to facilitate bot making"""
2020

2121
from .basepersistence import BasePersistence
22+
from .picklepersistence import PicklePersistence
2223
from .dispatcher import Dispatcher, DispatcherHandlerStop, run_async
2324
from .jobqueue import JobQueue, Job
2425
from .updater import Updater
@@ -44,4 +45,4 @@
4445
'MessageHandler', 'BaseFilter', 'Filters', 'RegexHandler', 'StringCommandHandler',
4546
'StringRegexHandler', 'TypeHandler', 'ConversationHandler',
4647
'PreCheckoutQueryHandler', 'ShippingQueryHandler', 'MessageQueue', 'DelayQueue',
47-
'DispatcherHandlerStop', 'run_async', 'BasePersistence')
48+
'DispatcherHandlerStop', 'run_async', 'BasePersistence', 'PicklePersistence')

‎telegram/ext/basepersistence.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
#!/usr/bin/env python
22
#
33
# A library that provides a Python interface to the Telegram Bot API
4-
# Copyright (C) 2015-2017
4+
# Copyright (C) 2015-2018
55
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
66
#
77
# This program is free software: you can redistribute it and/or modify
@@ -33,19 +33,19 @@ class BasePersistence(object):
3333
* :attr:`flush` will be called when the bot is shutdown, and must always be overwritten.
3434
3535
Attributes:
36-
store_user_data (:obj:`bool`): Whether user_data should be saved by this
36+
store_user_data (:obj:`bool`): Optional, Whether user_data should be saved by this
37+
persistence class.
38+
store_chat_data (:obj:`bool`): Optional. Whether user_data should be saved by this
3739
persistence class.
38-
store_chat_data (:obj:`bool`): Whether user_data should be saved by this
39-
persistence class
4040
4141
Args:
42-
store_user_data (:obj:`bool`): Whether user_data should be saved by this
43-
persistence class.
44-
store_chat_data (:obj:`bool`): Whether user_data should be saved by this
45-
persistence class
42+
store_user_data (:obj:`bool`, optional): Whether user_data should be saved by this
43+
persistence class. Default is ``True``.
44+
store_chat_data (:obj:`bool`, optional): Whether user_data should be saved by this
45+
persistence class. Default is ``True`` .
4646
"""
4747

48-
def __init__(self, store_user_data=False, store_chat_data=False):
48+
def __init__(self, store_user_data=True, store_chat_data=True):
4949
self.store_user_data = store_user_data
5050
self.store_chat_data = store_chat_data
5151

@@ -89,7 +89,7 @@ def update_conversations(self, name, conversations):
8989
9090
Args:
9191
name (:obj:`str`): The handlers name.
92-
conversation (:obj:'dict`): The :attr:`telegram.ext.ConversationHandler.conversations`
92+
conversations (:obj:'dict`): The :attr:`telegram.ext.ConversationHandler.conversations`
9393
dict to store.
9494
"""
9595
raise NotImplementedError

‎telegram/ext/picklepersistence.py‎

Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,229 @@
1+
#!/usr/bin/env python
2+
#
3+
# A library that provides a Python interface to the Telegram Bot API
4+
# Copyright (C) 2015-2018
5+
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
6+
#
7+
# This program is free software: you can redistribute it and/or modify
8+
# it under the terms of the GNU Lesser Public License as published by
9+
# the Free Software Foundation, either version 3 of the License, or
10+
# (at your option) any later version.
11+
#
12+
# This program is distributed in the hope that it will be useful,
13+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
14+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15+
# GNU Lesser Public License for more details.
16+
#
17+
# You should have received a copy of the GNU Lesser Public License
18+
# along with this program. If not, see [http://www.gnu.org/licenses/].
19+
"""This module contains the BasePersistence class."""
20+
import pickle
21+
from collections import defaultdict
22+
23+
24+
class PicklePersistence(object):
25+
"""Using python's builtin pickle for making you bot persistent.
26+
27+
Attributes:
28+
filename (:obj:`str`): The filename for storing the pickle files. When :attr:`single_file`
29+
is false this will be used as a prefix.
30+
store_user_data (:obj:`bool`): Optional. Whether user_data should be saved by this
31+
persistence class.
32+
store_chat_data (:obj:`bool`): Optional. Whether user_data should be saved by this
33+
persistence class.
34+
single_file (:obj:`bool`): Optional. When ``False`` will store 3 sperate files of
35+
`filename_user_data`, `filename_chat_data` and `filename_conversations`. Default is
36+
``True``.
37+
on_flush (:obj:`bool`): Optional. When ``True` will only save to file when :attr:`flush` is
38+
called and keep data in memory until that happens. When False will store data on any
39+
transaction. Default is ``False``.
40+
41+
Args:
42+
filename (:obj:`str`): The filename for storing the pickle files. When :attr:`single_file`
43+
is false this will be used as a prefix.
44+
store_user_data (:obj:`bool`, optional): Whether user_data should be saved by this
45+
persistence class. Default is ``True``.
46+
store_chat_data (:obj:`bool`, optional): Whether user_data should be saved by this
47+
persistence class. Default is ``True``.
48+
single_file (:obj:`bool`, optional): When ``False`` will store 3 sperate files of
49+
`filename_user_data`, `filename_chat_data` and `filename_conversations`. Default is
50+
``True``.
51+
on_flush (:obj:`bool`, optional): When ``True` will only save to file when :attr:`flush` is
52+
called and keep data in memory until that happens. When False will store data on any
53+
transaction. Default is ``False``.
54+
"""
55+
56+
def __init__(self, filename, store_user_data=True, store_chat_data=True, singe_file=True,
57+
on_flush=False):
58+
self.filename = filename
59+
self.store_user_data = store_user_data
60+
self.store_chat_data = store_chat_data
61+
self.single_file = singe_file
62+
self.on_flush = on_flush
63+
self.user_data = None
64+
self.chat_data = None
65+
self.conversations = None
66+
67+
def load_singlefile(self):
68+
try:
69+
filename = self.filename
70+
with open(self.filename, "rb") as f:
71+
all = pickle.load(f)
72+
self.user_data = defaultdict(dict, all['user_data'])
73+
self.chat_data = defaultdict(dict, all['chat_data'])
74+
self.conversations = all['conversations']
75+
except FileNotFoundError:
76+
self.conversations = {}
77+
self.user_data = defaultdict(dict)
78+
self.chat_data = defaultdict(dict)
79+
except pickle.UnpicklingError:
80+
raise TypeError("File {} does not contain valid pickle data".format(filename))
81+
82+
def load_file(self, filename):
83+
try:
84+
with open(filename, "rb") as f:
85+
return pickle.load(f)
86+
except FileNotFoundError:
87+
return None
88+
except pickle.UnpicklingError:
89+
raise TypeError("File {} does not contain valid pickle data".format(filename))
90+
91+
def dump_singlefile(self):
92+
with open(self.filename, "wb") as f:
93+
all = {'conversations': self.conversations, 'user_data': self.user_data,
94+
'chat_data': self.chat_data}
95+
pickle.dump(all, f)
96+
97+
def dump_file(self, filename, data):
98+
with open(filename, "wb") as f:
99+
pickle.dump(data, f)
100+
101+
def get_user_data(self):
102+
""""Returns the user_data from the pickle file if it exsists or an empty defaultdict.
103+
104+
Returns:
105+
:obj:'defaultdict`: The restored user data.
106+
"""
107+
if self.user_data:
108+
pass
109+
elif not self.single_file:
110+
filename = "{}_user_data".format(self.filename)
111+
data = self.load_file(filename)
112+
if not data:
113+
data = defaultdict(dict)
114+
else:
115+
data = defaultdict(dict, data)
116+
self.user_data = data
117+
else:
118+
self.load_singlefile()
119+
return self.user_data.copy()
120+
121+
def get_chat_data(self):
122+
""""Returns the chat_data from the pickle file if it exsists or an empty defaultdict.
123+
124+
Returns:
125+
:obj:'defaultdict`: The restored chat data.
126+
"""
127+
if self.chat_data:
128+
pass
129+
elif not self.single_file:
130+
filename = "{}_chat_data".format(self.filename)
131+
data = self.load_file(filename)
132+
if not data:
133+
data = defaultdict(dict)
134+
else:
135+
data = defaultdict(dict, data)
136+
self.chat_data = data
137+
else:
138+
self.load_singlefile()
139+
return self.chat_data.copy()
140+
141+
def get_conversations(self, name):
142+
""""Returns the conversations from the pickle file if it exsists or an empty defaultdict.
143+
144+
Args:
145+
name (:obj:`str`): The handlers name.
146+
147+
Returns:
148+
:obj:'dict`: The restored conversations for the handler.
149+
"""
150+
if self.conversations:
151+
pass
152+
elif not self.single_file:
153+
filename = "{}_conversations".format(self.filename)
154+
data = self.load_file(filename)
155+
if not data:
156+
data = {name: {}}
157+
self.conversations = data
158+
else:
159+
self.load_singlefile()
160+
return self.conversations.get(name, {}).copy()
161+
162+
def update_conversations(self, name, conversations):
163+
"""Will update the conversations for the given handler and depending on :attr:`on_flush`
164+
save the pickle file.
165+
166+
Args:
167+
name (:obj:`str`): The handlers name.
168+
conversations (:obj:'dict`): The :attr:`telegram.ext.ConversationHandler.conversations`
169+
dict to store.
170+
"""
171+
if self.conversations[name] == conversations:
172+
return
173+
self.conversations[name] = conversations
174+
if not self.on_flush:
175+
if not self.single_file:
176+
filename = "{}_conversations".format(self.filename)
177+
self.dump_file(filename, self.conversations)
178+
else:
179+
self.dump_singlefile()
180+
181+
def update_user_data(self, user_data):
182+
""""Will update the user_data (if changed) and depending on :attr:`on_flush` save the
183+
pickle file.
184+
185+
Args:
186+
user_data (:obj:'defaultdict`): The :attr:`telegram.ext.dispatcher.user_data`
187+
defaultdict to store.
188+
"""
189+
if self.user_data == user_data:
190+
return
191+
self.user_data = user_data
192+
if not self.on_flush:
193+
if not self.single_file:
194+
filename = "{}_user_data".format(self.filename)
195+
self.dump_file(filename, self.user_data)
196+
else:
197+
self.dump_singlefile()
198+
199+
def update_chat_data(self, chat_data):
200+
""""Will update the chat_data (if changed) and depending on :attr:`on_flush` save the
201+
pickle file.
202+
203+
Args:
204+
chat_data (:obj:'defaultdict`): The :attr:`telegram.ext.dispatcher.chat_data`
205+
defaultdict to store.
206+
"""
207+
if self.chat_data == chat_data:
208+
return
209+
self.chat_data = chat_data
210+
if not self.on_flush:
211+
if not self.single_file:
212+
filename = "{}_chat_data".format(self.filename)
213+
self.dump_file(filename, self.chat_data)
214+
else:
215+
self.dump_singlefile()
216+
217+
def flush(self):
218+
"""If :attr:`on_flush` is set to ``True``. Will save all data in memory to pickle file(s). If
219+
it's ``False`` will just pass.
220+
"""
221+
if not self.on_flush:
222+
pass
223+
else:
224+
if self.single_file:
225+
self.dump_singlefile()
226+
else:
227+
self.dump_file("{}_user_data".format(self.filename), self.user_data)
228+
self.dump_file("{}_chat_data".format(self.filename), self.chat_data)
229+
self.dump_file("{}_conversations".format(self.filename), self.conversations)

0 commit comments

Comments
 (0)