How to find exit code or reason when atexit callback is called in Python?

Question

I want to know if a Python script is terminating correctly or not. For this I am using atexit but the problem is that I do not know how to differentiate if atexit was called with sys.exit(0) or non zero or an exception.

Reasoning: if program ends properly, it will do nothing but if the program ends by an exception or returning an error code (exit status) different than zero I want to trigger some action.

In case you will wonder why I'm not using try/finally is because I want to add the same behaviour for a dozen of scripts that are importing a common module. Instead of modifying all of them, I want to add the atexit() hack to the module being imported and get this behaviour for free in all of them.

Answer 1

You can solve this using sys.excepthook and by monkey-patching sys.exit():

import atexit
import sys

class ExitHooks(object):
    def __init__(self):
        self.exit_code = None
        self.exception = None

    def hook(self):
        self._orig_exit = sys.exit
        sys.exit = self.exit
        sys.excepthook = self.exc_handler

    def exit(self, code=0):
        self.exit_code = code
        self._orig_exit(code)

    def exc_handler(self, exc_type, exc, *args):
        self.exception = exc

hooks = ExitHooks()
hooks.hook()

def foo():
    if hooks.exit_code is not None:
        print("death by sys.exit(%d)" % hooks.exit_code)
    elif hooks.exception is not None:
        print("death by exception: %s" % hooks.exception)
    else:
        print("natural death")
atexit.register(foo)

# test
sys.exit(1)