How can I register a function to be called only on *successful* exit of my Python program?

Question

I want to run a task when my Python program finishes, but only if it finishes successfully. As far as I know, using the atexit module means that my registered function will always be run at program termination, regardless of success. Is there a similar functionality to register a function so that it runs only on successful exit? Alternatively, is there a way for my exit function to detect whether the exit was normal or exceptional?

Here is some code that demonstrates the problem. It will print that the program succeeded, even when it has failed.

import atexit

def myexitfunc():
    print "Program succeeded!"

atexit.register(myexitfunc)

raise Exception("Program failed!")

Output:

$ python atexittest.py
Traceback (most recent call last):
  File "atexittest.py", line 8, in <module>
    raise Exception("Program failed!")
Exception: Program failed!
Program succeeded!

Answer 1

Out of the box, atexit is not quite suited for what you want to do: it's primarily used for resource cleanup at the very last moment, as things are shutting down and exiting. By analogy, it's the "finally" of a try/except, whereas what you want is the "else" of a try/except.

The simplest way I can think of is continuing to create a global flag which you set only when your script "succeeds"... and then have all the functions you attach to atexit check that flag, and do nothing unless it's been set.

Eg:

_success = False
def atsuccess(func, *args, **kwds):
    def wrapper():
        if _success:
            func(*args,**kwds)
    atexit(wrapper)

def set_success():
    global _success
    _success = True

# then call atsuccess() to attach your callbacks,
# and call set_success() before your script returns

One limitation is if you have any code which calls sys.exit(0) before setting the success flag. Such code should (probably) be refactored to return to the main function first, so that you call set_success and sys.exit in only one place. Failing that, you'll need add something like the following wrapper around the main entry point in your script:

try:
    main()
except SystemExit, err:
    if err.code == 0:
        set_success()
    raise

Answer 2

Wrap the body of your program in a with statement and define a corresponding context object that only performs your action when no exceptions have been raised. Something like:

class AtExit(object):
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_value is None:
            print "Success!"
        else:
            print "Failure!"

if __name__ == "__main__":
        with AtExit():
            print "Running"
#            raise Exception("Error")