yiskw note

機械学習やプログラミングについて気まぐれで書きます

【Python】プログラムの終了時に実行する関数を登録できるatexit

python

概要

今回はPythonの標準ライブラリであるatexitについて調べてみました。
非常に便利なライブラリですが、あまり日本語の情報がなかったので、こちらにメモを残しておきます。

atexitとは

docs.python.org

atexitはあまり知られていないかもしれませんが、プログラムの終了時に実行する関数を定義できる便利なライブラリです。
atexitを使用して登録した関数は、プログラムの正常終了時はもちろん、エラーが発生した場合やCtrl + Cで中断した場合にも実行させることが可能です。

ただし、プロセスをkillコマンドで止めた場合は、atexitで登録された関数を実行することができないので、
その場合は以下の記事が非常に参考になります。

qiita.com

使い方

atexit.registerで終了時に実行する関数を登録することができます。
atexit.registerは関数としても、デコレータとしても使用することができます。

試しに以下のようなコードを実行して、Ctrl + Cでプログラムを中止させてみます。

import atexit
import time


def main():
    time.sleep(100)


@atexit.register
def sendlog():
    print("\nProcess ended.")


if __name__ == "__main__":
    main()

実行結果

$ python sample.py  
^CTraceback (most recent call last):
  File "/Users/yuchi/Documents/projects/playground/sample.py", line 15, in <module>
    main()
  File "/Users/yuchi/Documents/projects/playground/sample.py", line 6, in main
    time.sleep(100)
KeyboardInterrupt

Process ended.

ちゃんと終了時にsendlog関数が実行され、"Process ended."と出力されました。

エラー時にも同様に動作します。

import atexit


def main():
    raise ValueError


@atexit.register
def sendlog():
    print("\nProcess ended.")


if __name__ == "__main__":
    main()

実行結果

$ python sample.py
Traceback (most recent call last):
  File "/Users/yuchi/Documents/projects/playground/sample.py", line 14, in <module>
    main()
  File "/Users/yuchi/Documents/projects/playground/sample.py", line 5, in main
    raise ValueError
ValueError

Process ended.

終了時に実行する関数に引数を追加することも可能です。
また、atexit.unregisterを使用すると、終了時に呼び出す登録をした関数の解除が可能です。

import atexit
import random


def print_value(n: int) -> None:
    print(f"\nProcess ended.")
    print(f"{n=}")


def main():
    n = random.randint(0, 100)

    atexit.register(print_value, n=n)

    if n < 50:
        raise ValueError

    # エラーが発生しなかった場合は、プログラム終了時にprint_valueを実行させない
    atexit.unregister(print_value)


if __name__ == "__main__":
    main()

実行結果 (エラー時)

$ python sample.py
Traceback (most recent call last):
  File "/Users/yuchi/Documents/projects/playground/sample.py", line 22, in <module>
    main()
  File "/Users/yuchi/Documents/projects/playground/sample.py", line 16, in main
    raise ValueError
ValueError

Process ended.
n=18

Reference