How to prevent a block of code from being interrupted by KeyboardInterrupt in Python?

Question

I'm writing a program that caches some results via the pickle module. What happens at the moment is that if I hit ctrl-c at while the dump operation is occurring, dump gets interrupted and the resulting file is corrupted (i.e. only partially written, so it cannot be loaded again.

Is there a way to make dump, or in general a block of code, uninterruptable? My current workaround looks something like this:

try:
  file = open(path, 'w')
  dump(obj, file)
  file.close()
except KeyboardInterrupt:
  file.close()
  file.open(path,'w')
  dump(obj, file)
  file.close()
  raise

It seems silly to restart the operation if it is interrupted, so I am searching for a way to defer the interrupt. How do I do this?

Answer 1

The following is a context manager that attaches a signal handler for SIGINT. If the context manager's signal handler is called, the signal is delayed by only passing the signal to the original handler when the context manager exits.

import signal
import logging

class DelayedKeyboardInterrupt:

    def __enter__(self):
        self.signal_received = False
        self.old_handler = signal.signal(signal.SIGINT, self.handler)
                
    def handler(self, sig, frame):
        self.signal_received = (sig, frame)
        logging.debug('SIGINT received. Delaying KeyboardInterrupt.')
    
    def __exit__(self, type, value, traceback):
        signal.signal(signal.SIGINT, self.old_handler)
        if self.signal_received:
            self.old_handler(*self.signal_received)

with DelayedKeyboardInterrupt():
    # stuff here will not be interrupted by SIGINT
    critical_code()