v0.0.1 (#1)

* [config] Initial file setup

* [new] Initial code and CI

* [docs] Updated readme

* [ci] adjust ci trigger

.github/FUNDING.yml

@@ -0,0 +1,3 @@
+# These are supported funding model platforms
+
+github: jeremyephron

.github/workflows/ci.yml

@@ -0,0 +1,40 @@
+name: CI
+
+on:
+  pull_request:
+    paths:
+      - "*"
+  workflow_dispatch:
+    paths:
+      - "*"
+  push:
+    paths:
+      - "setup.py"
+      - "pyterminate/*"
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: true
+      matrix:
+        os: ["ubuntu-latest", "macos-latest"]
+        python-version: ["3.6", "3.7", "3.8", "3.9", "3.10"]
+
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v2
+
+      - name: Setup python
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+          architecture: x64
+
+      - name: Install
+        run: |
+          pip install -e .
+          pip install -r requirements_test.txt
+
+      - name: Run tests
+        run: pytest

.github/workflows/python-publish.yml

@@ -0,0 +1,31 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [created]
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel twine
+    - name: Build and publish
+      env:
+        TWINE_USERNAME: '__token__'
+        TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+      run: |
+        python setup.py sdist bdist_wheel
+        twine upload dist/*

MANIFEST.in

@@ -0,0 +1 @@
+include LICENSE

README.md

@@ -1 +1,22 @@
-# pyterminate
+# pyterminate
+
+Reliably run cleanup upon program termination.
+
+## Quickstart
+
+```python3
+import signal
+
+import pyterminate
+
+@pyterminate.register(signals=(signal.SIGINT, signal.SIGTERM))
+def cleanup():
+    ...
+
+# or
+
+def cleanup(a, b):
+    ...
+
+pyterminate.register(cleanup, args=(None, 42))
+```

pyterminate/__init__.py

@@ -0,0 +1,75 @@
+"""
+File:
+-----
+
+"""
+
+from collections import defaultdict
+import atexit
+import signal
+import sys
+from types import FrameType
+from typing import Any, Callable, Dict, Optional, Tuple
+
+
+_registered_funcs = set()
+_func_to_wrapper = {}
+_signal_to_prev_handler = defaultdict(lambda: defaultdict(list))
+
+
+def register(
+    func: Optional[Callable] = None,
+    *,
+    args: tuple = tuple(),
+    kwargs: Optional[Dict[str, Any]] = None,
+    signals=(signal.SIGTERM,),
+    successful_exit: bool = False
+) -> Callable:
+    kwargs = kwargs or {}
+
+    def decorator(func: Callable) -> Callable:
+        return _register_impl(func, args, kwargs, signals, successful_exit)
+
+    return decorator(func) if func else decorator
+
+
+def unregister(func: Callable) -> None:
+    if func in _func_to_wrapper:
+        atexit.unregister(_func_to_wrapper[func])
+
+    _registered_funcs.remove(func)
+
+
+def _register_impl(
+    func: Callable,
+    args: tuple,
+    kwargs: Dict[str, Any],
+    signals: Tuple[int, ...],
+    successful_exit: bool,
+) -> Callable:
+    def exit_handler(*args: Any, **kwargs: Any) -> Any:
+        if func not in _registered_funcs:
+            return
+
+        _registered_funcs.remove(func)
+        return func(*args, **kwargs)
+
+    def signal_handler(sig: int, frame: FrameType):
+        signal.signal(sig, signal.SIG_IGN)
+
+        exit_handler(*args, **kwargs)
+        prev_handler = _signal_to_prev_handler[func][sig][-1]
+        if callable(prev_handler) and prev_handler != signal.default_int_handler:
+            prev_handler(sig, frame)
+
+        sys.exit(0 if successful_exit else sig)
+
+    for sig in signals:
+        _signal_to_prev_handler[func][sig].append(signal.signal(sig, signal_handler))
+
+    _registered_funcs.add(func)
+    _func_to_wrapper[func] = exit_handler
+
+    atexit.register(exit_handler, *args, **kwargs)
+
+    return func

requirements_test.txt

@@ -0,0 +1 @@
+pytest

setup.py

@@ -0,0 +1,23 @@
+import setuptools
+
+setuptools.setup(
+    name='pyterminate',
+    version='0.0.1',
+    url='https://github.com/jeremyephron/pyterminate',
+    author='Jeremy Ephron',
+    author_email='[email protected]',
+    description='Exit programs gracefully',
+    long_description=open('README.md').read(),
+    long_description_content_type='text/markdown',
+    packages=setuptools.find_packages(),
+    install_requires=[],
+    classifiers=[
+        'Development Status :: 3 - Alpha',
+        'Programming Language :: Python',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.9',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: MacOS',
+        'Operating System :: Unix',
+    ],
+)

tests/test_cleanup.py

@@ -0,0 +1,99 @@
+import multiprocessing as mp
+import os
+import signal
+import sys
+
+import pytest
+
+
+def setup_module():
+    mp.set_start_method('spawn')
+
+
+class ProcessUnderTest(mp.Process):
+    def __init__(self) -> None:
+        self.setup_is_done = mp.Event()
+        self.should_cleanup = mp.Event()
+
+        self._n_calls = mp.Value("i", 0)
+
+        self._pconn, self._cconn = mp.Pipe()
+
+        super().__init__(
+            target=self.run_program,
+            args=(
+                self._n_calls,
+                self.should_cleanup,
+                self.setup_is_done,
+                self._cconn
+            )
+        )
+
+    @property
+    def n_cleanup_calls(self):
+        return self._n_calls.value
+
+    def raise_exception(self, exc):
+        self._pconn.send(exc)
+
+    @staticmethod
+    def run_program(value, should_cleanup, setup_is_done, cconn) -> None:
+        import pyterminate
+
+        @pyterminate.register(signals=(signal.SIGINT, signal.SIGTERM))
+        def cleanup():
+            assert should_cleanup.wait(timeout=5)
+            value.value += 1
+
+        setup_is_done.set()
+
+        exc = cconn.recv()
+        if exc is not None:
+            raise exc
+
+        sys.exit(0)
+
+
+@pytest.fixture(scope='function')
+def proc() -> ProcessUnderTest:
+    return ProcessUnderTest()
+
+
+def test_normal_exit(proc: ProcessUnderTest) -> None:
+    proc.start()
+    proc.setup_is_done.wait()
+
+    proc.raise_exception(None)
+    proc.should_cleanup.set()
+
+    proc.join()
+
+    assert proc.exitcode == 0, proc.exitcode
+    assert proc.n_cleanup_calls == 1, proc.n_cleanup_calls
+
+
+def test_exception(proc: ProcessUnderTest) -> None:
+    proc.start()
+    proc.setup_is_done.wait()
+
+    proc.raise_exception(Exception("Something bad happened"))
+    proc.should_cleanup.set()
+
+    proc.join()
+
+    assert proc.exitcode == 1, proc.exitcode
+    assert proc.n_cleanup_calls == 1, proc.n_cleanup_calls
+
+
+@pytest.mark.parametrize('sig', [signal.SIGINT, signal.SIGTERM])
+def test_sigint(proc: ProcessUnderTest, sig: int) -> None:
+    proc.start()
+    proc.setup_is_done.wait()
+
+    os.kill(proc.pid, sig)
+    proc.should_cleanup.set()
+
+    proc.join()
+
+    assert proc.exitcode == sig, proc.exitcode
+    assert proc.n_cleanup_calls == 1, proc.n_cleanup_calls