Adjust persistence on exit behaviour (#1312)

* Adjust persistence of exit behaviour

* Fix binary operators in on_flush

* Fix docstring

* Add test

telegram/ext/dispatcher.py

@@ -395,6 +395,15 @@ class Dispatcher(object):
                 del self.handlers[group]
                 self.groups.remove(group)
 
+    def update_persistence(self):
+        """Update :attr:`user_data` and :attr:`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.
 

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 @@ class PicklePersistence(BasePersistence):
                 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 or self.chat_data or 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)

telegram/ext/updater.py

@@ -493,6 +493,8 @@ class Updater(object):
             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:

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 @@ class TestPickelPersistence(object):
         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)