Improved handlers

docs/features.md

@@ -463,3 +463,14 @@ def at_terminate():
     # do something before at_exit() if something went wrong
 ...
 ```
+
+Some of these behaviors can be disabled in a block of code using the `DisableSignals` context manager by passing it the signal identifiers.
+
+```python hl_lines="2"
+...
+with DisableSignals(SIGINT, SIGTERM) as _:
+    # do something
+    # if an interrupt or termination signal is raised, it will do nothing as
+    #  long as we are in this block of code
+...
+```

tests/test_handlers.py

@@ -12,23 +12,31 @@
 
 
 FILE = tmpf()
-SCRIPT = """from tinyscript import *
-def at_{}():
-    with open("{}", 'w+') as f:
-        f.write("{}")
-initialize(){}"""
+SCRIPT1 = """from tinyscript import *
+def at_{0}():
+    with open("{1}", 'w+') as f:
+        f.write("{2}")
+initialize()
+{3}"""
+SCRIPT2 = """from tinyscript import *
+def at_{0}():
+    with open("{1}", 'w+') as f:
+        f.write("{2}")
+initialize()
+with ts.DisableSignals({2}) as _:
+    {3}"""
 SIGNALS = {
     'interrupt': "SIGINT",
     'terminate': "SIGTERM",
 }
 TEXT = tmpf("handler-result", "txt")
 
 
-def exec_script(handler):
+def exec_script(handler, template):
     s = SIGNALS.get(handler)
-    s = ["\nos.kill(os.getpid(), signal.{})".format(s), ""][s is None]
+    s = ["os.kill(os.getpid(), signal.{})".format(s), ""][s is None]
     with open(FILE, 'w+') as f:
-        f.write(SCRIPT.format(handler.lower(), TEXT, handler.upper(), s))
+        f.write(template.format(handler.lower(), TEXT, handler.upper(), s))
     p = subprocess.Popen(["python{}".format(["2", "3"][PYTHON3]), FILE])
     p.wait()
     try:
@@ -41,13 +49,19 @@ def exec_script(handler):
 
 class TestHandlers(TestCase):
     def _test_handler(self, h):
-        self.assertEqual(exec_script(h), h.upper())
+        self.assertEqual(exec_script(h, SCRIPT1), h.upper())
 
     @classmethod
     def tearDownClass(self):
         remove(FILE)
         remove(TEXT)
 
+    def test_disable_handlers(self):
+        with DisableSignals(SIGINT) as _:
+            self.assertIsNone(exec_script("interrupt", SCRIPT2))
+            self.assertIsNone(exec_script("terminate", SCRIPT2))
+        self.assertRaises(ValueError, DisableSignals, 123456, fail=True)
+
     def test_exit_handler(self):
         self.assertIs(at_exit(), None)
         self._test_handler("exit")

tinyscript/handlers.py

@@ -3,15 +3,41 @@
 """Module for defining handlers and exit hook.
 
 """
-import logging
 import sys
-from signal import signal, SIGINT, SIGTERM
+from signal import getsignal, signal, SIGINT, SIGTERM
 
 
-__features__ = ["at_exit", "at_graceful_exit", "at_interrupt", "at_terminate"]
+__features__ = ["at_exit", "at_graceful_exit", "at_interrupt", "at_terminate",
+                "DisableSignals"]
 __all__ = ["_hooks"] + __features__
 
 
+class DisableSignals(object):
+    """
+    Context manager that disable signal handlers.
+
+    :param signals: list of signal identifiers
+    :param fail:    whether execution should fail or not when a bad signal ID is
+                     encountered
+    """
+    def __init__(self, *signals, **kwargs):
+        self.__handlers = {}
+        for s in signals:
+            try:
+                self.__handlers[s] = getsignal(s)
+            except ValueError as e:
+                if kwargs.get('fail', False):
+                    raise e
+
+    def __enter__(self):
+        for s in self.__handlers.keys():
+            signal(s, lambda *a, **kw: None)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        for s, h in self.__handlers.items():
+            signal(s, h)
+
+
 class ExitHooks(object):
     # inspired from: https://stackoverflow.com/questions/9741351/how-to-find-exi
     #                 t-code-or-reason-when-atexit-callback-is-called-in-python
@@ -26,7 +52,7 @@ def __init__(self):
     def exit(self, code=0):
         self.code = code
         self._orig_exit(code)
-    
+
     def quit(self, code=0):
         if self.state != "INTERRUPTED" or self._exit:
             self.exit(code)