Help us understand the problem. What is going on with this article?

Python で終了時に必ず何か実行したい

Pythonで終了時に必ず何か実行したい

途中でディレクトリを作るけど、終了時には残ってて欲しくない。
データベースに仮のデータを置くけど、終了時には残ってて欲しくない。
プログラムが終了したら正常な場合でも、死んだ場合でも、通知して欲しい。

このようなことを実現する方法です。

今回、やりたいこととして、以下のようなものを目指します。

  1. 正常な場合にはもちろん実行して欲しい
  2. Exceptionが発生しても実行して欲しい
  3. Ctrl-Cで止めても実行して欲しい
  4. killで止めても実行して欲しい
  5. Cleanup処理中は途中で止まらないで欲しい
  6. kill -9やSegmentaion Faultはあきらめる

いくつかの方法を試してみます。
結論を急ぐ方のために、先に結論を書いておきます。

先に結論

import sys
import time
import signal


def setup():
    print("!!!Set up!!!")


def cleanup():
    print("!!!Clean up!!!")
    # Cleanup処理いろいろ
    time.sleep(10)
    print("!!!Clean up Done!!!")


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


def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        # いろいろな処理
        time.sleep(60)

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_IGN)
        signal.signal(signal.SIGINT, signal.SIG_IGN)
        cleanup()
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        signal.signal(signal.SIGINT, signal.SIG_DFL)


if __name__ == "__main__":
    sys.exit(main())

解説編

準備

setup()/clean()関数とmainの部分を用意しておきます。

def setup():
    print("!!!Set up!!!")


def cleanup():
    print("!!!Clean up!!!")

def main():
    pass

if __name__ == "__main__":
    sys.exit(main())

Case1. try - finallyを使う

まず思いつくのはtry - finallyを使う方法です。

def main():
    setup()
    try:
        print("Do some jobs")

    finally:
        cleanup()

実行してみましょう。

!!!set up!!!
do some jobs
!!!Clean up!!!

Clean upされています。
それでは、途中でエラーが発生した場合はどうでしょうか。

def main():
    setup()
    try:
        print(1 / 0)

    finally:
        cleanup()

実行してみます。

!!!set up!!!
!!!Clean up!!!
Traceback (most recent call last):
  File "./try-finally.py", line 26, in <module>
    sys.exit(main())
  File "./try-finally.py", line 19, in main
    print(1 / 0)
ZeroDivisionError: division by zero

こちらも、Clean upされました。

では、Ctrl-Cで止めた場合はどうなるでしょうか。

def main():
    setup()
    try:
        time.sleep(60)

    finally:
        cleanup()

実行して、途中でCtrl-Cを押して止めてみます。

!!!Set up!!!
^C!!!Clean up!!!
Traceback (most recent call last):
  File "./try-finally.py", line 27, in <module>
    sys.exit(main())
  File "./try-finally.py", line 20, in main
    time.sleep(60)
KeyboardInterrupt

どうやらうまくいったようです。

では、killで殺した場合はどうでしょうか。

!!!Set up!!!
Terminated

これはいけません。これでは、ゴミファイルが残ってしまいます。

結論

try - finallyでは

  1. 正常な場合にはもちろん実行して欲しい ⇨ ○
  2. Exceptionが発生しても実行して欲しい ⇨ ○
  3. Ctrl-Cで止めても実行して欲しい ⇨ ○
  4. killで止めても実行して欲しい ⇨ ×

Case 2. atexitを使う

Pythonには、終了時に実行してくれるというatexitがあります。
これを試してみましょう。

def main():
    setup()
    atexit.register(cleanup)

    print("Do some jobs")

実行してみましょう。

!!!Set up!!!
Do some jobs
!!!Clean up!!!

うまくいきました。
次は、エラーを出してみます。

def main():
    setup()
    atexit.register(cleanup)

    print(1 / 0)

実行してみます。

!!!Set up!!!
Traceback (most recent call last):
  File "./try-finally.py", line 25, in <module>
    sys.exit(main())
  File "./try-finally.py", line 21, in main
    print(1 / 0)
ZeroDivisionError: division by zero
!!!Clean up!!!

こちらも、うまくいきました。
try - finallyの場合と比べるとClean upの処理がより後ろになっています。

それでは、Ctrl-Cで止めた場合はどうでしょうか。

def main():
    setup()
    atexit.register(cleanup)

    time.sleep(60)

実行して、途中でCtrl-Cで止めてみます。

!!!Set up!!!
^CTraceback (most recent call last):
  File "./try-finally.py", line 25, in <module>
    sys.exit(main())
  File "./try-finally.py", line 21, in main
    time.sleep(60)
KeyboardInterrupt
!!!Clean up!!!

いい感じです。

では、最後にkillで止めてみます。

!!!Set up!!!
Terminated

残念。try-finallyと同じようにこちらもダメなようです。

結論

atexitでは

  1. 正常な場合にはもちろん実行して欲しい ⇨ ○
  2. Exceptionが発生しても実行して欲しい ⇨ ○
  3. Ctrl-Cで止めても実行して欲しい ⇨ ○
  4. killで止めても実行して欲しい ⇨ ×

Case 3. signalを使う

killコマンドは、SIGTERMというsignalをプロセスに送ります。
また、Ctrl-CはSIGINTというsignalをプロセスに送ります。
これらのsignalをトラップしてみることにしましょう。

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


def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    signal.signal(signal.SIGINT, sig_handler)

    print("Do some jobs")

実行してみます。

!!!Set up!!!
Do some jobs

Clean upは実行されません。
外部からsignalが飛んできていないからですね。
あたりまえです。
Exceptionも同様なので、確認は省略します。

では、Ctrl-Cではどうなるか、確認してみます。

def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    signal.signal(signal.SIGINT, sig_handler)

    time.sleep(60)

実行して、途中でCtrl-Cで止めてみます。

!!!Set up!!!
^C!!!Clean up!!!

バッチリです。
KeyboardInterruptのExceptionは消えましたが、通常は問題にならないでしょう。
どうしても必要なら、sig_handlerの中でraise(KeyboardInterrupt())できます。

それでは、気になるkillですが、どうなるでしょうか。

!!!Set up!!!
!!!Clean up!!!

素晴らしい!
キチンとClean upされました。

結論

  1. 正常な場合にはもちろん実行して欲しい ⇨ ×
  2. Exceptionが発生しても実行して欲しい ⇨ ×
  3. Ctrl-Cで止めても実行して欲しい ⇨ ○
  4. killで止めても実行して欲しい ⇨ ○

ということで、try - finallyまたはatexitと、signalを組み合わせればよさそうです。

Case 4. try - finally と signalを組み合わせて使う

今回は、スコープが狭いという想定で、try - finallyとsignalを組み合わせてみることにします。

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


def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        print("do some jobs")

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        cleanup()

こんなふうになりました。
Ctrl-Cはtry - finallyでも動作するので、トラップするのはSIGTERMだけにしてみました。
tryを抜けた後は、cleanupは必要ないはずなので、finallyでSIGTERMをデフォルトに戻しておきます。

では、実行してみます。

!!!Set up!!!
do some jobs
!!!Clean up!!!

期待通りです。

次は、エラーを出してみます。

def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        print(1 / 0)

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        cleanup()


if __name__ == "__main__":
    sys.exit(main())

実行してみます。

!!!Set up!!!
!!!Clean up!!!
Traceback (most recent call last):
  File "./try-finally.py", line 33, in <module>
    sys.exit(main())
  File "./try-finally.py", line 26, in main
    print(1 / 0)
ZeroDivisionError: division by zero

Clean upされました。

次は、Ctrl-Cを確認します。

def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        time.sleep(60)

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        cleanup()

実行して、途中でCtrl-Cを押してみます。

!!!Set up!!!
^CClean up!!
Traceback (most recent call last):
  File "./try-finally.py", line 33, in <module>
    sys.exit(main())
  File "./try-finally.py", line 26, in main
    time.sleep(60)
KeyboardInterrupt

こちらも、Clean upされました。

さて、killはどうでしょうか。

!!!Set up!!!
!!!Clean up!!!
!!!Clean up!!!

えっ!
2回Clean upされてしまいました。

SIGTERMをトラップしたことで、Pythonが強制終了せずに、finallyのところも実行してくれたようです。
ということで、sig_handlerを次のように変更してみます。

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

handler内でcleanupを呼ぶのをやめて、終了だけするようにしました。

実行して、killしてみます。

!!!Set up!!!
!!!Clean up!!!

ようやく期待通りです。

結論

  1. 正常な場合にはもちろん実行して欲しい ⇨ ○
  2. Exceptionが発生しても実行して欲しい ⇨ ○
  3. Ctrl-Cで止めても実行して欲しい ⇨ ○
  4. killで止めても実行して欲しい ⇨ ○

Case 5. Cleanup処理中は止められないようにする

先ほどのやり方でも十分なのですが、Ctrl-Cを押したあと、すぐに止まらないと、Ctrl-Cを連打してしまうのが人情です。
killも何度も実行するかも知れません。

せっかくCleanup処理が動くようにしたのにそれでは台無しです。
ということで、Cleanup中は止められないようにします。

def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        time.sleep(60)

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_IGN)
        signal.signal(signal.SIGINT, signal.SIG_IGN)
        cleanup()
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        signal.signal(signal.SIGINT, signal.SIG_DFL)

cleanup()の前にCtrl-Cとkillのシグナルを無視するようにします。
cleanup()が終わったら、デフォルトに戻しておきます。
これで完璧です。

結論

  1. 正常な場合にはもちろん実行して欲しい ⇨ ○
  2. Exceptionが発生しても実行して欲しい ⇨ ○
  3. Ctrl-Cで止めても実行して欲しい ⇨ ○
  4. killで止めても実行して欲しい ⇨ ○
  5. 終了処理中は途中で止まらないで欲しい ⇨ ○

おまけ

try - finally / atexitの使い分け

今回の結果を見る限りtry-catchとatexitの違いは、スコープとタイミングの違いだけのようです。
setupがプログラムの途中で実行されるなら、そのスコープでtry - finallyを使う。
setupがプログラムの一番最初に実行するものなら、atexitを使う。
setupとか関係なく、最後にはいつでも実行して欲しいのなら、atexitを使う。
ということでしょうか。

あと、withを使う方法もあるので、興味のある方は試してみてください。

*本記事は @qualitia_cdevの中の一人、@hirachanさんに書いていただきました。

qualitia_cdev
QUALITIA CO., LTD. サービス企画部(aka 技術の無駄遣い部)が運営しているTechBlogです。常に新しい思考、新しい技術を追求し続けています。
https://www.qualitia.co.jp
qualitia
コミュニケーションの効率化とセキュリティの強化を支援するさまざまなメッセージング関連ソリューションを クラウド型サービス・ソフトウェアで提供しています。常に「クオリティの追求」への挑戦にこだわり、その企業活動とテクノロジーで社会に貢献することを目標としています。 
https://www.qualitia.co.jp/
Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away