Pythonで終了時に必ず何か実行したい (続編)

はじめに

3年ほど前に「Python で終了時に必ず何か実行したい」という記事を書きました。

3年も経つと、あちこちにこのコードのコピペが乱立してきました。
ふと、with句やdecoratorを使えばもうすこしきれいになるのではと思い、作ったものをpypiに公開してみました。

名前はFinalizerです。

pip install finalizer

で使えます。

ソースは

です。

使い方

使い方は簡単で、with Finalizer(掃除用の関数):のように指定します。
with句を抜けたときに掃除用関数が実行されます。
また、with句の実行中はCtrl-Cやkillコマンド(SIGTERM)はトラップされていますので、Ctrl-Cやkillコマンドで止めた場合も、掃除用関数は実行されます。

from time import sleep
from finalizer import Finalizer

def cleanup() -> None:
    print("cleanup start")
    sleep(3)
    print("cleanup end")

def task() -> None:
    print("task start")
    with Finalizer(cleanup):
        sleep(3)
        print("task end")

掃除用関数に引数が必要な場合は、

def cleanup(param1: str, param2: int) -> None:
    print("cleanup start", param1, param2)
    sleep(3)
    print("cleanup end")

def task2() -> None:
    print("task start")
    with Finalizer(cleanup, "test", param2=42):
        sleep(3)
        print("task end")

のように書くことができます。
with句はネストしても大丈夫です。

また、decoratorとしても指定できるので、次のように書くこともできます。

def cleanup() -> None:
    print("cleanup start")
    sleep(3)
    print("cleanup end")

@Finalizer(cleanup)
def task() -> None:
    print("task start")
    sleep(3)
    print("task end")

仕組みの開設

実際のソースコードより少しシンプルにしてありますが、骨子はこんな感じです。

from __future__ import annotations
from typing import Callable, Any
import signal
import sys
from contextlib import ContextDecorator


def sig_handler(signum, frame) -> None:
    sys.exit(1)


class Finalizer(ContextDecorator):
    def __init__(self, func: Callable[..., None], *args, **kw):
        self.__cleanup: Callable[..., None] = func
        self.__args: tuple[Any, ...] = args
        self.__kw: dict[str, Any] = kw
        self.__prev_sigterm: signal.Handlers = signal.SIG_DFL

    def __enter__(self):
        self.__prev_sigterm = signal.getsignal(signal.SIGTERM)
        signal.signal(signal.SIGTERM, sig_handler)

    def __exit__(self, exc_type, exc_value, traceback):
        cur_sigint = signal.getsignal(signal.SIGINT)
        signal.signal(signal.SIGTERM, signal.SIG_IGN)
        signal.signal(signal.SIGINT, signal.SIG_IGN)

        self.__cleanup(*self.__args, **self.__kw)

        signal.signal(signal.SIGTERM, self.__prev_sigterm)
        signal.signal(signal.SIGINT, cur_sigint)

with句実行時に__init__()と__enter__()が呼ばれます。
つまり、with句実行時に掃除用関数とその引数を保存しておきます。
同時に、現在のSIGTERMのハンドラも保存しておき、SIGTERMにsys.exit(1)を実行する関数をセットします。
SIGTERMにsys.exitをさせるようにしておかないと、Pythonで終了時に必ず何か実行したいで解説したtry～finallyの時と同様に、__exit__()が実行されずにそのまま終了してしまいます。

with句を抜けるときには__exit__()が実行されます。
掃除用関数実行中に再びCtrl-Cやkillコマンドで止められないようにSIGINTとSGTERMをIgnoreし、処理が終わったら元に戻します。

ちなみに、ContextDecratorを継承しておくことで、自動的にdecoratorにもなります。便利ですね。

おわりに

ごちゃごちゃせず、シンプルに書くことができるようになりました。
multi process環境でも動作します。

が、multi threadの環境では子threadでsignalを扱えないので、うまく動作しません。
もちろん、親threadでは使えるのですが、子threadは終了するまで待つしかありません。
結果として、掃除用関数は呼ばれるのですが、docker stopのような場合はkillされた後10秒後にkill -9されてしまうので、残念な結果になってしまいます。
何かいいアイデアがあれば、コメントやプルリクお待ちしています。