From 4455037ac9f72dd04be94f7c6437b1999d56e3a7 Mon Sep 17 00:00:00 2001
From: blep <baptiste.lepilleur@gmail.com>
Date: Sun, 27 Aug 2023 09:01:20 +0200
Subject: [PATCH] Added contextmanager for to automatically send WM_QUIT on
 CTRL+C or CTRL+Break.

---
 win32_window_monitor/__init__.py |  1 +
 win32_window_monitor/main.py     | 11 +----------
 win32_window_monitor/win32api.py | 32 +++++++++++++++++++++++++++++---
 3 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/win32_window_monitor/__init__.py b/win32_window_monitor/__init__.py
index b56a0b7..46ad222 100644
--- a/win32_window_monitor/__init__.py
+++ b/win32_window_monitor/__init__.py
@@ -12,4 +12,5 @@
     init_com,
     run_message_loop,
     post_quit_message,
+    post_quit_message_on_break_signal,
 )
\ No newline at end of file
diff --git a/win32_window_monitor/main.py b/win32_window_monitor/main.py
index be34d02..f09539d 100644
--- a/win32_window_monitor/main.py
+++ b/win32_window_monitor/main.py
@@ -57,7 +57,7 @@ def on_event(self, win_event_hook_handle, event_id: int, hwnd: wintypes.HWND,
 
 
 def main():
-    with init_com():
+    with init_com(), post_quit_message_on_break_signal():
         # Register hook callback for all relevant event types
         # Demonstrates that we can use a method as event hook callback without issue thanks
         # to ctypes.
@@ -65,15 +65,6 @@ def main():
         win_event_proc = WinEventProcType(event_logger.on_event)
         event_hook_handles = [set_win_event_hook(win_event_proc, et) for et in EVENT_TYPES.keys()]
 
-        # Install signal handler to exit the application when CTRL+C or CTRL+Break is pressed
-        def signal_handler(signum, frame):
-            # Send WM_QUIT message to exit the message loop started below
-            post_quit_message(0)
-
-        if platform.system() == 'Windows':
-            signal.signal(signal.SIGBREAK, signal_handler)
-        signal.signal(signal.SIGINT, signal_handler)
-
         # Run Windows message loop until WM_QUIT message is received (send by signal handlers above).
         # If you have a graphic UI, it is likely that your application already has a Windows message
         # loop that should be used instead.
diff --git a/win32_window_monitor/win32api.py b/win32_window_monitor/win32api.py
index c7bc07f..1dee286 100644
--- a/win32_window_monitor/win32api.py
+++ b/win32_window_monitor/win32api.py
@@ -1,10 +1,12 @@
+import contextlib
 import ctypes
-from ctypes import wintypes
 import logging
+import signal
+from ctypes import wintypes
 from typing import Optional
-from .ids import HookEvent, ObjectId
 from typing import Union
-import contextlib
+
+from .ids import HookEvent
 
 user32 = ctypes.windll.user32
 ole32 = ctypes.windll.ole32
@@ -174,3 +176,27 @@ def run_message_loop():
 
 def post_quit_message(exit_code: int = 0):
     PostQuitMessage(exit_code)
+
+
+@contextlib.contextmanager
+def post_quit_message_on_break_signal():
+    """Install signal handler to exit the application when CTRL+C or CTRL+Break is pressed.
+
+    Exiting the application is done by sending the WM_QUIT message (via post_quit_message),
+    which causes the Windows message loop of run_message_loop() that receives it to exit.
+
+    This is a contextmanager for use with the with statement.
+    """
+    def signal_handler_post_quit_message(signum, stack_frame):
+        post_quit_message()
+
+    old_break_handler = signal.getsignal(signal.SIGBREAK)
+    signal.signal(signal.SIGBREAK, signal_handler_post_quit_message)
+    old_int_handler = signal.getsignal(signal.SIGINT)
+    signal.signal(signal.SIGINT, signal_handler_post_quit_message)
+
+    yield
+
+    # Restore the old signal handlers on exit
+    signal.signal(signal.SIGINT, old_int_handler)
+    signal.signal(signal.SIGBREAK, old_break_handler)