gh-110771: Decompose run_forever() into parts (#110773)

Effectively introduce an unstable, private (really: protected) API for subclasses that want to override `run_forever()`.

4 files changed, 95 insertions, 31 deletions

diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py
index 956864e..0476de6 100644
--- a/Lib/asyncio/base_events.py
+++ b/Lib/asyncio/base_events.py
@@ -400,6 +400,8 @@ class BaseEventLoop(events.AbstractEventLoop):
         self._clock_resolution = time.get_clock_info('monotonic').resolution
         self._exception_handler = None
         self.set_debug(coroutines._is_debug_mode())
+        # The preserved state of async generator hooks.
+        self._old_agen_hooks = None
         # In debug mode, if the execution of a callback or a step of a task
         # exceed this duration in seconds, the slow callback/task is logged.
         self.slow_callback_duration = 0.1
@@ -601,29 +603,52 @@ class BaseEventLoop(events.AbstractEventLoop):
             raise RuntimeError(
                 'Cannot run the event loop while another loop is running')
 
-    def run_forever(self):
-        """Run until stop() is called."""
+    def _run_forever_setup(self):
+        """Prepare the run loop to process events.
+
+        This method exists so that custom custom event loop subclasses (e.g., event loops
+        that integrate a GUI event loop with Python's event loop) have access to all the
+        loop setup logic.
+        """
         self._check_closed()
         self._check_running()
         self._set_coroutine_origin_tracking(self._debug)
 
-        old_agen_hooks = sys.get_asyncgen_hooks()
-        try:
-            self._thread_id = threading.get_ident()
-            sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook,
-                                   finalizer=self._asyncgen_finalizer_hook)
+        self._old_agen_hooks = sys.get_asyncgen_hooks()
+        self._thread_id = threading.get_ident()
+        sys.set_asyncgen_hooks(
+            firstiter=self._asyncgen_firstiter_hook,
+            finalizer=self._asyncgen_finalizer_hook
+        )
+
+        events._set_running_loop(self)
+
+    def _run_forever_cleanup(self):
+        """Clean up after an event loop finishes the looping over events.
 
-            events._set_running_loop(self)
+        This method exists so that custom custom event loop subclasses (e.g., event loops
+        that integrate a GUI event loop with Python's event loop) have access to all the
+        loop cleanup logic.
+        """
+        self._stopping = False
+        self._thread_id = None
+        events._set_running_loop(None)
+        self._set_coroutine_origin_tracking(False)
+        # Restore any pre-existing async generator hooks.
+        if self._old_agen_hooks is not None:
+            sys.set_asyncgen_hooks(*self._old_agen_hooks)
+            self._old_agen_hooks = None
+
+    def run_forever(self):
+        """Run until stop() is called."""
+        try:
+            self._run_forever_setup()
             while True:
                 self._run_once()
                 if self._stopping:
                     break
         finally:
-            self._stopping = False
-            self._thread_id = None
-            events._set_running_loop(None)
-            self._set_coroutine_origin_tracking(False)
-            sys.set_asyncgen_hooks(*old_agen_hooks)
+            self._run_forever_cleanup()
 
     def run_until_complete(self, future):
         """Run until the Future is done.
diff --git a/Lib/asyncio/windows_events.py b/Lib/asyncio/windows_events.py
index 4a4c4be..b62ea75 100644
--- a/Lib/asyncio/windows_events.py
+++ b/Lib/asyncio/windows_events.py
@@ -314,24 +314,25 @@ class ProactorEventLoop(proactor_events.BaseProactorEventLoop):
             proactor = IocpProactor()
         super().__init__(proactor)
 
-    def run_forever(self):
-        try:
-            assert self._self_reading_future is None
-            self.call_soon(self._loop_self_reading)
-            super().run_forever()
-        finally:
-            if self._self_reading_future is not None:
-                ov = self._self_reading_future._ov
-                self._self_reading_future.cancel()
-                # self_reading_future was just cancelled so if it hasn't been
-                # finished yet, it never will be (it's possible that it has
-                # already finished and its callback is waiting in the queue,
-                # where it could still happen if the event loop is restarted).
-                # Unregister it otherwise IocpProactor.close will wait for it
-                # forever
-                if ov is not None:
-                    self._proactor._unregister(ov)
-                self._self_reading_future = None
+    def _run_forever_setup(self):
+        assert self._self_reading_future is None
+        self.call_soon(self._loop_self_reading)
+        super()._run_forever_setup()
+
+    def _run_forever_cleanup(self):
+        super()._run_forever_cleanup()
+        if self._self_reading_future is not None:
+            ov = self._self_reading_future._ov
+            self._self_reading_future.cancel()
+            # self_reading_future was just cancelled so if it hasn't been
+            # finished yet, it never will be (it's possible that it has
+            # already finished and its callback is waiting in the queue,
+            # where it could still happen if the event loop is restarted).
+            # Unregister it otherwise IocpProactor.close will wait for it
+            # forever
+            if ov is not None:
+                self._proactor._unregister(ov)
+            self._self_reading_future = None
 
     async def create_pipe_connection(self, protocol_factory, address):
         f = self._proactor.connect_pipe(address)
diff --git a/Lib/test/test_asyncio/test_base_events.py b/Lib/test/test_asyncio/test_base_events.py
index abcb6f5..c208097 100644
--- a/Lib/test/test_asyncio/test_base_events.py
+++ b/Lib/test/test_asyncio/test_base_events.py
@@ -922,6 +922,43 @@ class BaseEventLoopTests(test_utils.TestCase):
         self.loop.run_forever()
         self.loop._selector.select.assert_called_once_with(0)
 
+    def test_custom_run_forever_integration(self):
+        # Test that the run_forever_setup() and run_forever_cleanup() primitives
+        # can be used to implement a custom run_forever loop.
+        self.loop._process_events = mock.Mock()
+
+        count = 0
+
+        def callback():
+            nonlocal count
+            count += 1
+
+        self.loop.call_soon(callback)
+
+        # Set up the custom event loop
+        self.loop._run_forever_setup()
+
+        # Confirm the loop has been started
+        self.assertEqual(asyncio.get_running_loop(), self.loop)
+        self.assertTrue(self.loop.is_running())
+
+        # Our custom "event loop" just iterates 10 times before exiting.
+        for i in range(10):
+            self.loop._run_once()
+
+        # Clean up the event loop
+        self.loop._run_forever_cleanup()
+
+        # Confirm the loop has been cleaned up
+        with self.assertRaises(RuntimeError):
+            asyncio.get_running_loop()
+        self.assertFalse(self.loop.is_running())
+
+        # Confirm the loop actually did run, processing events 10 times,
+        # and invoking the callback once.
+        self.assertEqual(self.loop._process_events.call_count, 10)
+        self.assertEqual(count, 1)
+
     async def leave_unfinalized_asyncgen(self):
         # Create an async generator, iterate it partially, and leave it
         # to be garbage collected.
diff --git a/Misc/NEWS.d/next/Library/2023-10-13-06-47-20.gh-issue-110771.opwdlc.rst b/Misc/NEWS.d/next/Library/2023-10-13-06-47-20.gh-issue-110771.opwdlc.rst
new file mode 100644
index 0000000..a22f8a0
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2023-10-13-06-47-20.gh-issue-110771.opwdlc.rst
@@ -0,0 +1 @@
+Expose the setup and cleanup portions of ``asyncio.run_forever()`` as the standalone methods ``asyncio.run_forever_setup()`` and ``asyncio.run_forever_cleanup()``. This allows for tighter integration with GUI event loops.