How to invoke something whenever python is finished?

I make a temporary directory, but I don't want it remains after finished.
I put the temporary data into the database, but I don't want them remain after finished.
I want to get a notification whenever python did completed whether it's normal or dead.

I will explain the magic to accomplish the preceding tasks.

First, let me define what we want.

1. Invoke something when it has finished normally.
2. Invoke something when the Exception has been raised.
3. Invoke something after stopping by Ctrl-C.
4. Invoke something after stopping by the kill command.
5. I don't want it to stop while invoking the 'Cleanup'.
6. Give up kill -9 or Segmentation Fault.

Ok, so let's try several methods.

But before that, for those who are so busy, the following is my conclusion.

Conclusion

import sys
import time
import signal


def setup():
    print("!!!Set up!!!")


def cleanup():
    print("!!!Clean up!!!")
    # some cleanup tasks
    time.sleep(10)
    print("!!!Clean up Done!!!")


def sig_handler(signum, frame) -> None:
    sys.exit(1)


def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        # some tasks
        time.sleep(60)

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_IGN)
        signal.signal(signal.SIGINT, signal.SIG_IGN)
        cleanup()
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        signal.signal(signal.SIGINT, signal.SIG_DFL)


if __name__ == "__main__":
    sys.exit(main())

Commentary

Preparations

Let's prepare setup()/clean() functions and main part of program.

def setup():
    print("!!!Set up!!!")


def cleanup():
    print("!!!Clean up!!!")

def main():
    pass

if __name__ == "__main__":
    sys.exit(main())

Case 1. Use try - finally

At first, I come up with to use try - finally.

def main():
    setup()
    try:
        print("Do some jobs")

    finally:
        cleanup()

Let's test it!

!!!set up!!!
do some jobs
!!!Clean up!!!

Ok, clean up is invoked.
So, what will happen when the error is occurred?

def main():
    setup()
    try:
        print(1 / 0)

    finally:
        cleanup()

Let's test.

!!!set up!!!
!!!Clean up!!!
Traceback (most recent call last):
  File "./try-finally.py", line 26, in <module>
    sys.exit(main())
  File "./try-finally.py", line 19, in main
    print(1 / 0)
ZeroDivisionError: division by zero

This is also ok, "Cleanup" was invoked.

So, what will happen when stopped by Ctrl-C?

def main():
    setup()
    try:
        time.sleep(60)

    finally:
        cleanup()

Run the code and press Ctrl-C to stop it.

!!!Set up!!!
^C!!!Clean up!!!
Traceback (most recent call last):
  File "./try-finally.py", line 27, in <module>
    sys.exit(main())
  File "./try-finally.py", line 20, in main
    time.sleep(60)
KeyboardInterrupt

Oh, it seems good!

Ok, so let's stop it by kill command.

!!!Set up!!!
Terminated

OMG! This is too bad! Cleanup was not invoked. There must remain much garbage.

Conclusion

In the case to use try - finally:

1. Invoke something when it has finished normally. ==> OK
2. Invoke something when the Exception has been raised. ==> OK
3. Invoke something after stopping by Ctrl-C. ==> OK
4. Invoke something after stopping by the kill command. ==> FAILED

Case 2. Use atexit

Python has a useful function named atexit which calls a function when it finishes.
Let's try with this.

def main():
    setup()
    atexit.register(cleanup)

    print("Do some jobs")

Test it!

!!!Set up!!!
Do some jobs
!!!Clean up!!!

Ok, good.
Next, make an error.

def main():
    setup()
    atexit.register(cleanup)

    print(1 / 0)

Test it.

!!!Set up!!!
Traceback (most recent call last):
  File "./try-finally.py", line 25, in <module>
    sys.exit(main())
  File "./try-finally.py", line 21, in main
    print(1 / 0)
ZeroDivisionError: division by zero
!!!Clean up!!!

This is also ok.
Clean up method is called later than the case try - finally.

So, what will happen when stopping by Ctrl-C?

def main():
    setup()
    atexit.register(cleanup)

    time.sleep(60)

Run the code, and press Ctrl-C to stop it.

!!!Set up!!!
^CTraceback (most recent call last):
  File "./try-finally.py", line 25, in <module>
    sys.exit(main())
  File "./try-finally.py", line 21, in main
    time.sleep(60)
KeyboardInterrupt
!!!Clean up!!!

Ok, sounds nice.

So, let's stop by kill command.

!!!Set up!!!
Terminated

Oh, so disappointed. This is as same as try-finally.

Conclusion

In the case to use atexit:

1. Invoke something when it has finished normally. ==> OK
2. Invoke something when the Exception has been raised. ==> OK
3. Invoke something after stopping by Ctrl-C. ==> OK
4. Invoke something after stopping by the kill command. ==> FAILED

Case 3. Use signal

The kill command sends SIGTERM signal to the process.
Ctrl-C sends SIGINT signal to the process.
Let's trap these signals.

def sig_handler(signum, frame) -> None:
    cleanup()
    sys.exit(1)


def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    signal.signal(signal.SIGINT, sig_handler)

    print("Do some jobs")

Test it!

!!!Set up!!!
Do some jobs

Of course, clean up has not been invoked.
Nobody sends signal in the normal path.
The exception case is also the same.

So, what will happen when pressing Ctrl-C?

def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    signal.signal(signal.SIGINT, sig_handler)

    time.sleep(60)

Run the code and press Ctrl-C to stop it.

!!!Set up!!!
^C!!!Clean up!!!

Good!
KeyboardInterrupt Exception is disappeared, but nobody may care.
If you need it, you can raise(KeyboardInterrupt()) in sig_handler.

So, how about kill command? Let's see.

!!!Set up!!!
!!!Clean up!!!

Wow! Beautiful!!
Clean up has been invoked.

Conclusion

1. Invoke something when it has finished normally. ==> FAILED
2. Invoke something when the Exception has been raised. ==> FAILED
3. Invoke something after stopping by Ctrl-C. ==> OK
4. Invoke something after stopping by the kill command. ==> OK

Ok, I think we can mix try - finally or atexit, and signal.

Case 4. Use mixed try - finally and signal

This time let's assume the scope is narrow.
I will try to use mixing try - finally and signal.

def sig_handler(signum, frame) -> None:
    cleanup()
    sys.exit(1)


def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        print("do some jobs")

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        cleanup()

I wrote like above.
I trapped the only SIGTERM, because Ctrl-C case can work well with try-finally.
Outside of the try clause, cleanup is already not necessary, I set the SIGTERM to default behavior in the finally clause.

Ok, so test it.

!!!Set up!!!
do some jobs
!!!Clean up!!!

Ok, as I expected.

Next, error is occurred.

def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        print(1 / 0)

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        cleanup()


if __name__ == "__main__":
    sys.exit(main())

Test it.

!!!Set up!!!
!!!Clean up!!!
Traceback (most recent call last):
  File "./try-finally.py", line 33, in <module>
    sys.exit(main())
  File "./try-finally.py", line 26, in main
    print(1 / 0)
ZeroDivisionError: division by zero

Cleanup has been called.

Next, test Ctrl-C.

def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        time.sleep(60)

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        cleanup()

Run the code and press Ctrl-C.

!!!Set up!!!
^CClean up!!
Traceback (most recent call last):
  File "./try-finally.py", line 33, in <module>
    sys.exit(main())
  File "./try-finally.py", line 26, in main
    time.sleep(60)
KeyboardInterrupt

This is also ok! Cleanup has been called.

So, how about kill command?

!!!Set up!!!
!!!Clean up!!!
!!!Clean up!!!

Eh!!
Cleanup was called twice!

As a result of trapping SIGTERM, python has not terminated abnormally and invoked finally clause correctly.

So, let's change the sig_handler function like this.

def sig_handler(signum, frame) -> None:
    sys.exit(1)

This time, sig_handler does not call cleanup function but only exit().

Run the code and kill it.

!!!Set up!!!
!!!Clean up!!!

Ok! Finally, this is what I expected!

Conclusion

1. Invoke something when it has finished normally. ==> OK
2. Invoke something when the Exception has been raised. ==> OK
3. Invoke something after stopping by Ctrl-C. ==> OK
4. Invoke something after stopping by the kill command. ==> OK

Case 5. Avoid stopping during Cleanup

Case 4 is almost enough, but if the cleanup function takes long time, you might press Ctrl-C again and again.
Also, you might run kill command again and again.

Even if you implement cleanup method, there might still remain garbage.
So, let's prevent stopping during cleanup process.

def main():
    setup()
    signal.signal(signal.SIGTERM, sig_handler)
    try:
        time.sleep(60)

    finally:
        signal.signal(signal.SIGTERM, signal.SIG_IGN)
        signal.signal(signal.SIGINT, signal.SIG_IGN)
        cleanup()
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        signal.signal(signal.SIGINT, signal.SIG_DFL)

Before cleanup() function, ignore the signals from Ctrl-C and the kill command.
And make it default after cleanup().
This is perfect!

Conclusion

1. Invoke something when it has finished normally. ==> OK
2. Invoke something when the Exception has been raised. ==> OK
3. Invoke something after stopping by Ctrl-C. ==> OK
4. Invoke something after stopping by the kill command. ==> OK
5. I don't want it to stop whlie invoking the 'Cleanup'. ==> OK

Extras

What's the differece between try - finally and atexit?

I think the differences between try-finally and atexit are only the scope and timing.
If the setup is called in the middle of the program, use try-finally in that scope.
If the setup is called at the beginning of the program, use atexit.
If you don't care about setup, use atexit anyway.

Also, there is a way to use with.
If you have interests, let's try it!