2024 0825 atexit

Discovered this week that Python has an exit handler module called atexit. Exit handlers can be registered with atexit.register(handler, *args, **kwargs), and they are run when the script exits in the oppossite order they are registered.

This can be particularly useful in programs that perform work in stages and need to clean up each stage by the end of the process whether it executed normally or with an error.

Instead of

try:
    stage_one()
    try:
        stage_two()
    finally:
        cleanup_stage_two()
finally:
    cleanup_stage_one()
#... etc

you can do

atexit.register(cleanup_stage_one)
stage_one()
# If the script fails here, cleanup_stage_one will run before exit
atexit.register(cleanup_stage_two)
stage_two()
# If the script fails here, cleanup_stage_two will run, then cleanup_stage_one, then it will exit
#... etc

You need to be careful with signal handling, however; the documentation says “The functions registered via this module are not called when the program is killed by a signal not handled by Python, when a Python fatal internal error is detected, or when os._exit() is called.” To have the registered functions run even when being killed by a signal, for instance with kill $your_pid, you must explicitly handle each signal you care about.

def signal_handler(signum, frame):
    """Handle signal a signal.

    atexit.register(...) does not handle signals UNLESS the signal is caught by a handler.
    So we need to a define a handler that just logs and calls sys.exit(),
    and let atexit do any cleanup.
    """
    logger.warning(f"Received signal {signum}, cleaning up and exiting.")
    # The normal convention is for programs exiting due to a sent signal
    # to return 128 plus the signal number.
    sys.exit(128 + signum)

for sig in [
    signal.SIGHUP,
    signal.SIGINT,
    signal.SIGQUIT,
    signal.SIGPIPE,
    signal.SIGTERM,
]:
    signal.signal(sig, signal_handler)

This requires knowing which signals you need to care about, which is a little arcane. My normal list is the one above.

It’s similar in conception to, but more limited than, Go’s defer.

I used this to write a script¹ that builds this site on a cloud VM. Each stage depends on the previous stage, but the whole thing needs to be cleaned up when the script exits whether the build succeeds or fails, so atexit was perfect.