Gracefully Exiting Python Context Managers on Ctrl+C

In this post I will show you how you can exit gracefully from within a context manager in your Python application.

Suppose that we provide a context manager that can be used as a session to perform certain tasks. When the context manager is closed, there is some cleanup work to be done. We want this to happen even when the user interrupts the program with Ctrl+C key.

This is how we can achieve that:

import sys
import time
from signal import signal, SIGINT

class Session:

    def __enter__(self):
        signal(SIGINT, self._sigint_handler)


    def __exit__(self, type, value, traceback):
        print('Exiting session...')


    def _do_cleanup(self):
        print('Cleaning up...')

    def _sigint_handler(self, signal_received, frame):
        print('Ctrl + C handler called')

        self.__exit__(None, None, None)

if __name__ == '__main__':
    with Session():
        print('Session started')

If you run the program above, the program will sleep and block for a long time. When you press Ctrl+C, the handler we registered (_sigint_handler) will be called instead of a KeyboardInterrupt error being raised. In this method we manually call the context manager's __exit__ method, passing a value of None for all 3 arguments. This is what is actually passed when __exit__ is routinely called when the context manager is finished.

Notice that after manually calling __exit__, we are also manually terminating the program with sys.exit(0). This is necessary, since omitting this will cause our program to keep running. Pressing Ctrl+C will not terminate the program.



comments powered by Disqus