How can I register a function that will only be called when * successfully * exits my Python program?

Question

I want to run a task when my Python program exits, but only if it successfully completes. As far as I know, using the atexit module means that my registered function will always work when the program terminates, regardless of success. Is there a similar function for registering a function so that it only works on successful completion? Alternatively, is there a way for my exit function to determine if the output was normal or exceptional?

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

import atexit def myexitfunc(): print "Program succeeded!" atexit.register(myexitfunc) raise Exception("Program failed!") 

Conclusion:

 $ 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 not quite suitable for what you want to do: first of all, it is used to clean resources at the very last moment when everything stops and exits. By analogy, this is "finally" try / except, whereas you want an "else" try / except.

The easiest way I can think of is to create a global flag that you set only when your script "succeeds" ... and then all the functions that you attach to atexit check this flag and do nothing if it doesn't it was found.

For instance:

 _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 code that calls sys.exit(0) before setting the success flag. Such code should (probably) be reorganized to return to the main function, so you call set_success and sys.exit only one place. Otherwise, you will need to add something like the following shell around the main entry point in the 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 an appropriate context object that only performs your actions when no exceptions have been made. Sort of:

 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")