From 4c97b328f1ff228257de5761b1b3748bd0a7232e Mon Sep 17 00:00:00 2001 From: cyc60 Date: Wed, 28 Jun 2023 15:45:33 +0300 Subject: [PATCH] Transform InterruptHandler to context manager --- sw_utils/common.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/sw_utils/common.py b/sw_utils/common.py index 534b7c2..841cec8 100644 --- a/sw_utils/common.py +++ b/sw_utils/common.py @@ -7,16 +7,30 @@ class InterruptHandler: """ Tracks SIGINT and SIGTERM signals. + Usage: + with InterruptHandler() as interrupt_handler: + while not interrupt_handler.exit: + ... """ exit = False - def __init__(self) -> None: + def __enter__(self) -> 'InterruptHandler': signal.signal(signal.SIGINT, self.exit_gracefully) signal.signal(signal.SIGTERM, self.exit_gracefully) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + signal.signal(signal.SIGINT, self.exit_default) + signal.signal(signal.SIGTERM, self.exit_default) - # noinspection PyUnusedLocal def exit_gracefully(self, signum: int, *args, **kwargs) -> None: # pylint: disable=unused-argument + if self.exit: + raise KeyboardInterrupt logger.info('Received interrupt signal %s, exiting...', signum) self.exit = True + + def exit_default(self, signum: int, *args, **kwargs) -> None: + # pylint: disable=unused-argument + raise KeyboardInterrupt