Source code for labml.utils.delayed_keyboard_interrupt

import signal

from labml import logger
from labml.logger import Text


class DelayedKeyboardInterrupt:
    """
    When this is used in a ``with`` block it will capture keyboard interrupts and
    fire them at the end of the ``with`` block after all the code completes execution.
    """

    def __init__(self):
        self.signal_received = None

    def __enter__(self):
        self.signal_received = None
        # Start capturing
        self.old_handler = signal.signal(signal.SIGINT, self.handler)

    def handler(self, sig, frame):
        # Pass second interrupt without delaying
        if self.signal_received is not None:
            self.old_handler(*self.signal_received)
            return

        # Store the interrupt signal for later
        self.signal_received = (sig, frame)
        logger.log([('\nSIGINT received. Delaying KeyboardInterrupt.',
                     Text.danger)])

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Reset handler
        signal.signal(signal.SIGINT, self.old_handler)

        # Pass on any captured interrupt signals
        if self.signal_received is not None:
            self.old_handler(*self.signal_received)