[pygobject/benzea/gio-asyncio: 8/9] overrides: Add EventLoop integration points into overrides

[Benjamin Berg <bberg src gnome org>]

commit 04d4596f9f1797dac39653e705170d15835263c3
Author: Benjamin Berg <bberg redhat com>
Date:   Tue Dec 28 00:03:44 2021 +0100

    overrides: Add EventLoop integration points into overrides
    
    This adds EventLoop integration in order to mark the event loop as
    running when a main context iterating API is called. Note that for the
    simple APIs to only do one iteration the EventLoop is paused (i.e. it
    will not dispatch). No one should do this, but it might happen when e.g.
    porting and this should create a well-defined behaviour.

 gi/_ossighelper.py   | 35 +++++++++++++++++++++++++++++++++++
 gi/events.py         |  4 ++++
 gi/overrides/GLib.py |  7 ++++---
 gi/overrides/Gio.py  |  4 ++--
 gi/overrides/Gtk.py  | 21 +++++++++++++++++----
 5 files changed, 62 insertions(+), 9 deletions(-)
---
diff --git a/gi/_ossighelper.py b/gi/_ossighelper.py
index fd1b4499..543bbe12 100644
--- a/gi/_ossighelper.py
+++ b/gi/_ossighelper.py
@@ -18,6 +18,7 @@ from __future__ import print_function
 import os
 import socket
 import signal
+import asyncio
 import threading
 from contextlib import closing, contextmanager
 
@@ -238,3 +239,37 @@ def register_sigint_fallback(callback):
             signal.default_int_handler(signal.SIGINT, None)
         else:
             _callback_stack.pop()
+
+
+class DummyEventLoop():
+    @classmethod
+    @contextmanager
+    def paused(cls):
+        yield
+
+    @classmethod
+    @contextmanager
+    def running(cls, quit_func):
+        with wakeup_on_signal():
+            yield
+
+
+def get_event_loop(ctx):
+    """Return the correct GLibEventLoop or a dummy that just registers the
+    signal wakeup mechanism."""
+
+    # Try to use the running loop. If there is none, get the policy and
+    # try getting one in the hope that this will give us an event loop for the
+    # correct context.
+    loop = asyncio._get_running_loop()
+    if loop is None:
+        try:
+            loop = asyncio.get_event_loop_policy().get_event_loop_for_context(ctx)
+        except:
+            pass
+
+    if loop and hasattr(loop, '_context'):
+        if ctx is not None and hash(loop._context) == hash(ctx):
+            return loop
+
+    return DummyEventLoop
diff --git a/gi/events.py b/gi/events.py
index 96156331..95c2ac8d 100644
--- a/gi/events.py
+++ b/gi/events.py
@@ -302,6 +302,10 @@ class GLibEventLoopPolicy(asyncio.AbstractEventLoopPolicy):
             raise RuntimeError('There is no main context set for thread %r.'
                                % threading.current_thread().name)
 
+        return self.get_event_loop_for_context(ctx)
+
+    def get_event_loop_for_context(self, ctx):
+        """Get the event loop for a specific context."""
         # Note: We cannot attach it to ctx, as getting the default will always
         #       return a new python wrapper. But, we can use hash() as that returns
         #       the pointer to the C structure.
diff --git a/gi/overrides/GLib.py b/gi/overrides/GLib.py
index 78d309b6..cd03eb03 100644
--- a/gi/overrides/GLib.py
+++ b/gi/overrides/GLib.py
@@ -23,7 +23,7 @@ import warnings
 import sys
 import socket
 
-from .._ossighelper import wakeup_on_signal, register_sigint_fallback
+from .._ossighelper import register_sigint_fallback, get_event_loop
 from ..module import get_introspection_module
 from .._gi import (variant_type_from_string, source_new,
                    source_set_callback, io_channel_read)
@@ -493,7 +493,7 @@ class MainLoop(GLib.MainLoop):
 
     def run(self):
         with register_sigint_fallback(self.quit):
-            with wakeup_on_signal():
+            with get_event_loop(self.get_context()).running(self.quit):
                 super(MainLoop, self).run()
 
 
@@ -504,7 +504,8 @@ __all__.append('MainLoop')
 class MainContext(GLib.MainContext):
     # Backwards compatible API with default value
     def iteration(self, may_block=True):
-        return super(MainContext, self).iteration(may_block)
+        with get_event_loop(self).paused():
+            return super(MainContext, self).iteration(may_block)
 
 
 MainContext = override(MainContext)
diff --git a/gi/overrides/Gio.py b/gi/overrides/Gio.py
index c807fe0b..d7daae2d 100644
--- a/gi/overrides/Gio.py
+++ b/gi/overrides/Gio.py
@@ -20,7 +20,7 @@
 
 import warnings
 
-from .._ossighelper import wakeup_on_signal, register_sigint_fallback
+from .._ossighelper import register_sigint_fallback, get_event_loop
 from ..overrides import override, deprecated_init, wrap_list_store_sort_func
 from ..module import get_introspection_module
 from gi import PyGIWarning
@@ -38,7 +38,7 @@ class Application(Gio.Application):
 
     def run(self, *args, **kwargs):
         with register_sigint_fallback(self.quit):
-            with wakeup_on_signal():
+            with get_event_loop(GLib.MainContext.default()).running(self.quit):
                 return Gio.Application.run(self, *args, **kwargs)
 
 
diff --git a/gi/overrides/Gtk.py b/gi/overrides/Gtk.py
index f53ec6bd..bd99ea75 100644
--- a/gi/overrides/Gtk.py
+++ b/gi/overrides/Gtk.py
@@ -22,8 +22,8 @@
 import sys
 import warnings
 
-from gi.repository import GObject
-from .._ossighelper import wakeup_on_signal, register_sigint_fallback
+from gi.repository import GObject, GLib
+from .._ossighelper import register_sigint_fallback, get_event_loop
 from .._gtktemplate import Template, _extract_handler_and_args
 from ..overrides import (override, strip_boolean_result, deprecated_init,
                          wrap_list_store_sort_func)
@@ -581,7 +581,7 @@ class Dialog(Gtk.Dialog, Container):
 
         def run(self, *args, **kwargs):
             with register_sigint_fallback(self.destroy):
-                with wakeup_on_signal():
+                with get_event_loop(GLib.MainContext.default()).running(self.destroy):
                     return Gtk.Dialog.run(self, *args, **kwargs)
 
         action_area = property(lambda dialog: dialog.get_action_area())
@@ -1683,9 +1683,22 @@ if GTK2 or GTK3:
     @override(Gtk.main)
     def main(*args, **kwargs):
         with register_sigint_fallback(Gtk.main_quit):
-            with wakeup_on_signal():
+            with get_event_loop(GLib.MainContext.default()).running(Gtk.main_quit):
                 return _Gtk_main(*args, **kwargs)
 
+    _Gtk_main_iteration = Gtk.main_iteration
+    _Gtk_main_iteration_do = Gtk.main_iteration_do
+
+    @override(Gtk.main_iteration)
+    def main_iteration():
+        with get_event_loop(GLib.MainContext.default()).paused():
+            return _Gtk_main_iteration()
+
+    @override(Gtk.main_iteration)
+    def main_iteration_do(blocking):
+        with get_event_loop(GLib.MainContext.default()).paused():
+            return _Gtk_main_iteration_do(blocking)
+
 
 if GTK2 or GTK3:
     stock_lookup = strip_boolean_result(Gtk.stock_lookup)