v0.0.3 (#3)

* [update] Ignore all requested signals during call

- Ignores signals that a function was registered to during its execution
  to ensure completion. They are restored after execution.
- Added explanation of decorator and some nuances of this behavior in
  the README
- Added tests for the new decorator
- Fix ordering issue in in unregister unit test

* [test] Added test for duplicate registration

Ensure function only gets called once if registered multiple times.

README.md

@@ -10,6 +10,7 @@ Reliably run cleanup code upon program termination.
 - [Why does this exist?](#why-does-this-exist)
 - [What can it do?](#what-can-it-do)
 - [Quickstart](#quickstart)
+- [Tips and tricks](#tips-and-tricks)
 
 ## Why does this exist?
 
@@ -52,11 +53,10 @@ This packages does or allows the following behavior:
 
 - Allows functions to be unregistered: `pyterminate.unregister(func)`
 
-- Ignore multiple repeated signals to allow the registered functions to
-  complete
-  - However, it can be canceled upon receipt of another signal. Desired
-    behavior could vary application to application, but this feels appropriate
-    for the most common known use cases.
+- Ignore requested signals while registered function is executing, ensuring
+  that it is not interrupted.
+  - It's important to note that `SIGKILL` and calls to `os._exit()` cannot be
+    ignored.
 
 ## Quickstart
 
@@ -84,3 +84,51 @@ def cleanup(*args, **kwargs):
 
 pyterminate.register(cleanup, ...)
 ```
+
+## Tips and tricks
+
+### Multiprocessing start method
+
+When starting processes with Python's
+[`multiprocessing`](https://docs.python.org/3/library/multiprocessing.html)
+module, the `fork` method will fail to call registered functions on exit, since
+the process is ended with `os._exit()` internally, which bypasses all cleanup
+and immediately kills the process.
+
+One way of getting around this are using the `"spawn"` start method if that
+is acceptable for your application. Another method is to register your function
+to a user-defined signal, and wrap your process code in try-except block,
+raising the user-defined signal at the end. `pyterminate` provides this
+functionality in the form of the `exit_with_signal` decorator, which simply
+wraps the decorated function in a try-finally block, and raises the given
+signal. Example usage:
+
+```python3
+import multiprocessing as mp
+import signal
+
+import pyterminate
+
+
+@pyterminate.exit_with_signal(signal.SIGUSR1)
+def run_process():
+
+    @pyterminate.register(signals=[signal.SIGUSR1, signal.SIGINT, signal.SIGTERM])
+    def cleanup():
+        ...
+
+    ...
+
+
+if __name__ == "__main__"
+    mp.set_start_method("fork")
+
+    proc = mp.Process(target=run_process)
+    proc.start()
+
+    try:
+        proc.join(timeout=300)
+    except TimeoutError:
+        proc.terminate()
+        proc.join()
+```

pyterminate/__init__.py

@@ -8,6 +8,8 @@
 
 from collections import defaultdict
 import atexit
+import functools
+import os
 import signal
 import sys
 from types import FrameType
@@ -94,11 +96,9 @@ def _register_impl(
     This function is the internal implementation of registration, and should
     not be called by a user, who should called register() instead.
 
-    Idempotent handlers are created for both atexit and signal handling.
-    The currently handled signal is ignored during the signal handler to allow
-    for the registered functions to complete when potentially receiving
-    multiple repeated signals. However, it can be canceled upon receipt of
-    another signal.
+    Idempotent handlers are created for both atexit and signal handling. All
+    requested signals are ignored while registered function is executing, and
+    are restored afterward.
 
     Args:
         func: Function to register.
@@ -115,16 +115,21 @@ def _register_impl(
 
     """
 
-    def exit_handler(*args: Any, **kwargs: Any) -> Any:
+    def exit_handler(*args: Any, **kwargs: Any) -> None:
         if func not in _registered_funcs:
             return
 
+        prev_handlers = {}
+        for sig in signals:
+            prev_handlers[sig] = signal.signal(sig, signal.SIG_IGN)
+
         _registered_funcs.remove(func)
-        return func(*args, **kwargs)
+        func(*args, **kwargs)
 
-    def signal_handler(sig: int, frame: Optional[FrameType]) -> None:
-        signal.signal(sig, signal.SIG_IGN)
+        for sig, handler in prev_handlers.items():
+            signal.signal(sig, handler)
 
+    def signal_handler(sig: int, frame: Optional[FrameType]) -> None:
         exit_handler(*args, **kwargs)
 
         if _signal_to_prev_handler[func][sig]:
@@ -153,3 +158,29 @@ def signal_handler(sig: int, frame: Optional[FrameType]) -> None:
     atexit.register(exit_handler, *args, **kwargs)
 
     return func
+
+
+def exit_with_signal(signum: int) -> Callable:
+    """
+    Creates a decorator for raising a given signal on exit.
+
+    Args:
+        signum: The signal to raise on exit.
+
+    Returns:
+        decorator: Decorator that executes a function and raises a signal
+            afterward.
+
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> None:
+            try:
+                func(*args, **kwargs)
+            finally:
+                os.kill(os.getpid(), signum)
+
+        return wrapper
+
+    return decorator

setup.py

@@ -2,7 +2,7 @@
 
 setuptools.setup(
     name='pyterminate',
-    version='0.0.2',
+    version='0.0.3',
     url='https://github.com/jeremyephron/pyterminate',
     author='Jeremy Ephron',
     author_email='[email protected]',

tests/test_cleanup.py

@@ -6,6 +6,8 @@
 
 import pytest
 
+import pyterminate
+
 ProgramType = Callable[[Dict[str, Any], mp.Value, mp.Event, mp.Event, mp.Pipe], None]
 
 
@@ -17,9 +19,9 @@ class ProcessUnderTest(mp.Process):
     def __init__(self, program: ProgramType) -> None:
         self.setup_is_done = mp.Event()
         self.should_cleanup = mp.Event()
-        self.register_kwargs = {"signals": (signal.SIGINT, signal.SIGTERM)}
+        self.register_kwargs = {'signals': (signal.SIGINT, signal.SIGTERM)}
 
-        self._n_calls = mp.Value("i", 0)
+        self._n_calls = mp.Value('i', 0)
 
         self._pconn, self._cconn = mp.Pipe()
 
@@ -101,10 +103,9 @@ def unregister_program(
     def cleanup(a=1, b=0):
         value.value += (a + b)
 
-    setup_is_done.set()
-
     pyterminate.unregister(cleanup)
 
+    setup_is_done.set()
     assert should_cleanup.wait(timeout=5)
 
     sys.exit(0)
@@ -142,6 +143,35 @@ def cleanup3(a=1, b=0):
     sys.exit(0)
 
 
+def duplicate_register_program(
+    register_kwargs,
+    value,
+    should_cleanup,
+    setup_is_done,
+    cconn
+) -> None:
+    import pyterminate
+
+    signal.signal(signal.SIGTERM, lambda *_: sys.exit(66))
+
+    @pyterminate.register(**register_kwargs)
+    def cleanup(a=1, b=0):
+        value.value += (a + b)
+
+    pyterminate.register(cleanup, **register_kwargs)
+    pyterminate.register(cleanup, **register_kwargs)
+
+    setup_is_done.set()
+    assert should_cleanup.wait(timeout=5)
+
+    sys.exit(0)
+
+
+@pyterminate.exit_with_signal(signal.SIGUSR1)
+def simple_program_exit_with_signal(*args: Any, **kwargs: Any) -> None:
+    simple_program(*args, **kwargs)
+
+
 @pytest.fixture(scope='function')
 def proc(request) -> ProcessUnderTest:
     return ProcessUnderTest(getattr(request, "param", simple_program))
@@ -301,3 +331,59 @@ def test_multiple_register(proc: ProcessUnderTest) -> None:
 
     assert proc.exitcode == 66
     assert proc.n_cleanup_calls == 2
+
+
+@pytest.mark.parametrize(
+    'proc', [simple_program_exit_with_signal], indirect=True
+)
+def test_exit_with_signal(proc: ProcessUnderTest) -> None:
+    """Tests that the exit_with_signal decorator works on exceptions."""
+
+    proc.register_kwargs.update(signals=(signal.SIGUSR1,))
+
+    proc.start()
+    proc.setup_is_done.wait()
+
+    proc.raise_exception(None)
+    proc.should_cleanup.set()
+
+    proc.join()
+
+    assert proc.exitcode == signal.SIGUSR1
+    assert proc.n_cleanup_calls == 1
+
+
+@pytest.mark.parametrize(
+    'proc', [simple_program_exit_with_signal], indirect=True
+)
+def test_exit_with_signal_on_exc(proc: ProcessUnderTest) -> None:
+    """Tests that the exit_with_signal decorator works on exceptions."""
+
+    proc.register_kwargs.update(signals=(signal.SIGUSR1,))
+
+    proc.start()
+    proc.setup_is_done.wait()
+
+    proc.raise_exception(Exception())
+    proc.should_cleanup.set()
+
+    proc.join()
+
+    assert proc.exitcode == signal.SIGUSR1
+    assert proc.n_cleanup_calls == 1
+
+
+@pytest.mark.parametrize('proc', [duplicate_register_program], indirect=True)
+def test_duplicate_register(proc: ProcessUnderTest) -> None:
+    """
+    Tests that the function is only called once with duplicate registrations.
+
+    """
+
+    proc.start()
+    proc.setup_is_done.wait()
+    proc.should_cleanup.set()
+    proc.join()
+
+    assert proc.exitcode == 0
+    assert proc.n_cleanup_calls == 1