30

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.

sorin
  • 161,544
  • 178
  • 535
  • 806

2 Answers2

24

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)
Pynchia
  • 10,996
  • 5
  • 34
  • 43
Niklas B.
  • 92,950
  • 18
  • 194
  • 224
  • @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. – Niklas B. Mar 16 '12 at 17:32
  • @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. – Niklas B. Mar 16 '12 at 17:32
  • 1
    @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). – Niklas B. Mar 16 '12 at 17:45
  • 4
    This seems to make the traceback stdout/err disappear. Is there a way to keep the traceback, and just add the print statements? – e9t Jun 29 '15 at 09:32
  • @e9t 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 – Kyle Hannon Mar 27 '20 at 21:28
1

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)
Maxi
  • 421
  • 4
  • 4