From 45bb776f607c6a42d9306ce0c10da654a0ed4bde Mon Sep 17 00:00:00 2001 From: Bibo-Joshi Date: Sun, 16 Dec 2018 13:43:02 +0100 Subject: [PATCH 1/4] Adjust persistence of exit behaviour --- telegram/ext/dispatcher.py | 9 +++++++++ telegram/ext/picklepersistence.py | 24 ++++++++++++------------ telegram/ext/updater.py | 2 ++ 3 files changed, 23 insertions(+), 12 deletions(-) diff --git a/telegram/ext/dispatcher.py b/telegram/ext/dispatcher.py index 2ee47f11d4f..66d34bce614 100644 --- a/telegram/ext/dispatcher.py +++ b/telegram/ext/dispatcher.py @@ -395,6 +395,15 @@ def remove_handler(self, handler, group=DEFAULT_GROUP): del self.handlers[group] self.groups.remove(group) + def update_persistence(self): + """Update ``user_data`` and ``chat_data`` in :attr:`persistence`. + """ + if self.persistence: + for chat_id in self.chat_data: + self.persistence.update_chat_data(chat_id, self.chat_data[chat_id]) + for user_id in self.user_data: + self.persistence.update_user_data(user_id, self.user_data[user_id]) + def add_error_handler(self, callback): """Registers an error handler in the Dispatcher. diff --git a/telegram/ext/picklepersistence.py b/telegram/ext/picklepersistence.py index ed3c06cf900..0f2ecef8b92 100644 --- a/telegram/ext/picklepersistence.py +++ b/telegram/ext/picklepersistence.py @@ -36,9 +36,9 @@ class PicklePersistence(BasePersistence): single_file (:obj:`bool`): Optional. When ``False`` will store 3 sperate files of `filename_user_data`, `filename_chat_data` and `filename_conversations`. Default is ``True``. - on_flush (:obj:`bool`): Optional. When ``True`` will only save to file when :meth:`flush` - is called and keep data in memory until that happens. When False will store data on any - transaction. Default is ``False``. + on_flush (:obj:`bool`, optional): When ``True`` will only save to file when :meth:`flush` + is called and keep data in memory until that happens. When ``False`` will store data + on any transaction *and* on call fo :meth:`flush`. Default is ``False``. Args: filename (:obj:`str`): The filename for storing the pickle files. When :attr:`single_file` @@ -51,8 +51,8 @@ class PicklePersistence(BasePersistence): `filename_user_data`, `filename_chat_data` and `filename_conversations`. Default is ``True``. on_flush (:obj:`bool`, optional): When ``True`` will only save to file when :meth:`flush` - is called and keep data in memory until that happens. When False will store data on any - transaction. Default is ``False``. + is called and keep data in memory until that happens. When ``False`` will store data + on any transaction *and* on call fo :meth:`flush`. Default is ``False``. """ def __init__(self, filename, store_user_data=True, store_chat_data=True, singe_file=True, @@ -221,15 +221,15 @@ def update_chat_data(self, chat_id, data): self.dump_singlefile() def flush(self): - """If :attr:`on_flush` is set to ``True``. Will save all data in memory to pickle file(s). If - it's ``False`` will just pass. + """ Will save all data in memory to pickle file(s). """ - if not self.on_flush: - pass - else: - if self.single_file: + if self.single_file: + if self.user_data and self.chat_data and self.conversations: self.dump_singlefile() - else: + else: + if self.user_data: self.dump_file("{}_user_data".format(self.filename), self.user_data) + if self.chat_data: self.dump_file("{}_chat_data".format(self.filename), self.chat_data) + if self.conversations: self.dump_file("{}_conversations".format(self.filename), self.conversations) diff --git a/telegram/ext/updater.py b/telegram/ext/updater.py index 3f050e04255..638b2036463 100644 --- a/telegram/ext/updater.py +++ b/telegram/ext/updater.py @@ -493,6 +493,8 @@ def signal_handler(self, signum, frame): self.logger.info('Received signal {} ({}), stopping...'.format( signum, get_signal_name(signum))) if self.persistence: + # Update user_data and chat_data before flushing + self.dispatcher.update_persistence() self.persistence.flush() self.stop() if self.user_sig_handler: From 51202367ca5b345e93bb3a40e6cc5905c850bb50 Mon Sep 17 00:00:00 2001 From: Bibo-Joshi Date: Sun, 13 Jan 2019 12:43:55 +0100 Subject: [PATCH 2/4] Fix binary operators in on_flush --- telegram/ext/picklepersistence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/telegram/ext/picklepersistence.py b/telegram/ext/picklepersistence.py index 0f2ecef8b92..d15ade20a79 100644 --- a/telegram/ext/picklepersistence.py +++ b/telegram/ext/picklepersistence.py @@ -224,7 +224,7 @@ def flush(self): """ Will save all data in memory to pickle file(s). """ if self.single_file: - if self.user_data and self.chat_data and self.conversations: + if self.user_data or self.chat_data or self.conversations: self.dump_singlefile() else: if self.user_data: From 918c7e11950452e40135b232b36fab53cdc5434d Mon Sep 17 00:00:00 2001 From: Pieter Schutz Date: Wed, 13 Feb 2019 22:45:25 +0100 Subject: [PATCH 3/4] Fix docstring --- telegram/ext/dispatcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/telegram/ext/dispatcher.py b/telegram/ext/dispatcher.py index 66d34bce614..bec5b681d3c 100644 --- a/telegram/ext/dispatcher.py +++ b/telegram/ext/dispatcher.py @@ -396,7 +396,7 @@ def remove_handler(self, handler, group=DEFAULT_GROUP): self.groups.remove(group) def update_persistence(self): - """Update ``user_data`` and ``chat_data`` in :attr:`persistence`. + """Update :attr:`user_data` and :attr:`chat_data` in :attr:`persistence`. """ if self.persistence: for chat_id in self.chat_data: From 7c64fd5a615255547902554a39862e97e02ec63d Mon Sep 17 00:00:00 2001 From: Pieter Schutz Date: Wed, 13 Feb 2019 22:46:24 +0100 Subject: [PATCH 4/4] Add test --- tests/test_persistence.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_persistence.py b/tests/test_persistence.py index c43c160580b..f3fb627f691 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -16,6 +16,8 @@ # # You should have received a copy of the GNU Lesser Public License # along with this program. If not, see [http://www.gnu.org/licenses/]. +import signal + from telegram.utils.helpers import enocde_conversations_to_json try: @@ -539,6 +541,24 @@ def second(bot, update, user_data, chat_data): dp.add_handler(h2) dp.process_update(update) + def test_flush_on_stop(self, bot, update, pickle_persistence, good_pickle_files): + u = Updater(bot=bot, persistence=pickle_persistence) + dp = u.dispatcher + u.running = True + dp.user_data[4242424242]['my_test'] = 'Working!' + dp.chat_data[-4242424242]['my_test2'] = 'Working2!' + u.signal_handler(signal.SIGINT, None) + del (dp) + del (u) + del (pickle_persistence) + pickle_persistence_2 = PicklePersistence(filename='pickletest', + store_user_data=True, + store_chat_data=True, + singe_file=False, + on_flush=False) + assert pickle_persistence_2.get_user_data()[4242424242]['my_test'] == 'Working!' + assert pickle_persistence_2.get_chat_data()[-4242424242]['my_test2'] == 'Working2!' + def test_with_conversationHandler(self, dp, update, good_pickle_files, pickle_persistence): dp.persistence = pickle_persistence NEXT, NEXT2 = range(2)