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.

Rockbound

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)

Heterologous

@sorin: It works fine for me. I wrapped the hook functions in a class (which doesn't change functionality) and added the sys.exit(1). It prints sys.exit(1) for me on Python 2.7. – Heterologous

@sorin: Did you maybe forget a global exit_code inside patched_exit? I forgot it at first, but edited it in right after posting the answer. – Heterologous

@sorin: I only switched to a class to avoid the global variables, the mechanism is exactly the same. If you check the old version, it should work just as well (at least it does for me). – Heterologous

This seems to make the traceback stdout/err disappear. Is there a way to keep the traceback, and just add the print statements? – Weasner

@Weasner I'm late to the game here, but just started using this and ran into your issue. You need to capture the original excepthook and then call it. Basically what this solution does for sys.exit, but for sys.excepthook – Anergy

Answer 2

This adaption of Niklas. B.`s code shows the stack trace.

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
        self._orig_exc_handler = self.exc_handler
        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
        self._orig_exc_handler(self, exc_type, exc, *args)

def exit_handler():
    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")

hooks = ExitHooks()
hooks.hook()
atexit.register(exit_handler)

# test
sys.exit(1)

Conclude