Added user defined function for updater's signalHandler (python-telegram-bot#512)

d-qoi · jh0ker · commit 5b14b134dcbb · 2017-03-26T14:36:13.000+02:00
* Added user defined function for updater's signalHandler

* Added test_userSignal to test_updater

* Added test_userSignal to test_updater
Fixing paren
diff --git a/AUTHORS.rst b/AUTHORS.rst
@@ -14,6 +14,7 @@ The following wonderful people contributed directly or indirectly to this projec
 - `Anton Tagunov <https://github.com/anton-tagunov>`_
 - `Balduro <https://github.com/Balduro>`_
 - `bimmlerd <https://github.com/bimmlerd>`_
+- `d-qoi <https://github.com/d-qoi>`
 - `daimajia <https://github.com/daimajia>`_
 - `Eli Gao <https://github.com/eligao>`_
 - `ErgoZ Riftbit Vaper <https://github.com/ergoz>`_
diff --git a/telegram/ext/updater.py b/telegram/ext/updater.py
@@ -62,6 +62,9 @@ class Updater(object):
         bot (Optional[Bot]): A pre-initialized bot instance. If a pre-initizlied bot is used, it is
             the user's responsibility to create it using a `Request` instance with a large enough
             connection pool.
+        user_sig_handler (Optional[function]): Takes ``signum, frame`` as positional arguments.
+            This will be called when a signal is received, defaults are (SIGINT, SIGTERM, SIGABRT)
+            setable with Updater.idle(stop_signals=(signals))
         request_kwargs (Optional[dict]): Keyword args to control the creation of a request object
             (ignored if `bot` argument is used).
 
@@ -71,7 +74,14 @@ class Updater(object):
     """
     _request = None
 
-    def __init__(self, token=None, base_url=None, workers=4, bot=None, request_kwargs=None):
+    def __init__(self,
+                 token=None,
+                 base_url=None,
+                 workers=4,
+                 bot=None,
+                 user_sig_handler=None,
+                 request_kwargs=None):
+
         if (token is None) and (bot is None):
             raise ValueError('`token` or `bot` must be passed')
         if (token is not None) and (bot is not None):
@@ -92,6 +102,7 @@ def __init__(self, token=None, base_url=None, workers=4, bot=None, request_kwarg
                 request_kwargs['con_pool_size'] = workers + 4
             self._request = Request(**request_kwargs)
             self.bot = Bot(token, base_url, request=self._request)
+        self.user_sig_handler = user_sig_handler
         self.update_queue = Queue()
         self.job_queue = JobQueue(self.bot)
         self.__exception_event = Event()
@@ -426,6 +437,8 @@ def signal_handler(self, signum, frame):
         self.is_idle = False
         if self.running:
             self.stop()
+            if self.user_sig_handler:
+                self.user_sig_handler(signum, frame)
         else:
             self.logger.warning('Exiting immediately!')
             import os
diff --git a/tests/test_updater.py b/tests/test_updater.py
@@ -790,6 +790,23 @@ def test_idle(self):
         sleep(1)
         self.assertFalse(self.updater.running)
 
+    def test_userSignal(self):
+        self._setup_updater('Test7', messages=0)
+
+        tempVar = {'a': 0}
+
+        def userSignalInc(signum, frame):
+            tempVar['a'] = 1
+
+        self.updater.user_sig_handler = userSignalInc
+        self.updater.start_polling(poll_interval=0.01)
+        Thread(target=self.signalsender).start()
+        self.updater.idle()
+        # If we get this far, idle() ran through
+        sleep(1)
+        self.assertFalse(self.updater.running)
+        self.assertTrue(tempVar['a'] != 0)
+
     def test_createBot(self):
         self.updater = Updater('123:abcd')
         self.assertIsNotNone(self.updater.bot)
