Add exception handler to webloop (#1452)

src/pyodide-py/pyodide/webloop.py

@@ -1,6 +1,8 @@
 import asyncio
 import time
 import contextvars
+import sys
+import traceback
 
 
 from typing import Callable
@@ -26,6 +28,8 @@ class WebLoop(asyncio.AbstractEventLoop):
     def __init__(self):
         self._task_factory = None
         asyncio._set_running_loop(self)
+        self._exception_handler = None
+        self._current_handle = None
 
     def get_debug(self):
         return False
@@ -111,7 +115,7 @@ class WebLoop(asyncio.AbstractEventLoop):
         delay: float,
         callback: Callable,
         *args,
-        context: contextvars.Context = None
+        context: contextvars.Context = None,
     ):
         """Arrange for a callback to be called at a given time.
 
@@ -144,7 +148,7 @@ class WebLoop(asyncio.AbstractEventLoop):
         when: float,
         callback: Callable,
         *args,
-        context: contextvars.Context = None
+        context: contextvars.Context = None,
     ):
         """Like ``call_later()``, but uses an absolute time.
 
@@ -238,6 +242,134 @@ class WebLoop(asyncio.AbstractEventLoop):
         """
         return self._task_factory
 
+    def get_exception_handler(self):
+        """Return an exception handler, or None if the default one is in use."""
+        return self._exception_handler
+
+    def set_exception_handler(self, handler):
+        """Set handler as the new event loop exception handler.
+
+        If handler is None, the default exception handler will be set.
+
+        If handler is a callable object, it should have a signature matching
+        '(loop, context)', where 'loop' will be a reference to the active event
+        loop, 'context' will be a dict object (see `call_exception_handler()`
+        documentation for details about context).
+        """
+        if handler is not None and not callable(handler):
+            raise TypeError(
+                f"A callable object or None is expected, " f"got {handler!r}"
+            )
+        self._exception_handler = handler
+
+    def default_exception_handler(self, context):
+        """Default exception handler.
+
+        This is called when an exception occurs and no exception handler is set,
+        and can be called by a custom exception handler that wants to defer to
+        the default behavior. This default handler logs the error message and
+        other context-dependent information.
+
+
+        In debug mode, a truncated stack trace is also appended showing where
+        the given object (e.g. a handle or future or task) was created, if any.
+        The context parameter has the same meaning as in
+        `call_exception_handler()`.
+        """
+        message = context.get("message")
+        if not message:
+            message = "Unhandled exception in event loop"
+
+        exception = context.get("exception")
+        if exception is not None:
+            exc_info = (type(exception), exception, exception.__traceback__)
+        else:
+            exc_info = False
+
+        if (
+            "source_traceback" not in context
+            and self._current_handle is not None
+            and self._current_handle._source_traceback
+        ):
+            context["handle_traceback"] = self._current_handle._source_traceback
+
+        log_lines = [message]
+        for key in sorted(context):
+            if key in {"message", "exception"}:
+                continue
+            value = context[key]
+            if key == "source_traceback":
+                tb = "".join(traceback.format_list(value))
+                value = "Object created at (most recent call last):\n"
+                value += tb.rstrip()
+            elif key == "handle_traceback":
+                tb = "".join(traceback.format_list(value))
+                value = "Handle created at (most recent call last):\n"
+                value += tb.rstrip()
+            else:
+                value = repr(value)
+            log_lines.append(f"{key}: {value}")
+
+        print("\n".join(log_lines), file=sys.stderr)
+
+    def call_exception_handler(self, context):
+        """Call the current event loop's exception handler.
+        The context argument is a dict containing the following keys:
+        - 'message': Error message;
+        - 'exception' (optional): Exception object;
+        - 'future' (optional): Future instance;
+        - 'task' (optional): Task instance;
+        - 'handle' (optional): Handle instance;
+        - 'protocol' (optional): Protocol instance;
+        - 'transport' (optional): Transport instance;
+        - 'socket' (optional): Socket instance;
+        - 'asyncgen' (optional): Asynchronous generator that caused
+                                 the exception.
+        New keys maybe introduced in the future.
+        Note: do not overload this method in an event loop subclass.
+        For custom exception handling, use the
+        `set_exception_handler()` method.
+        """
+        if self._exception_handler is None:
+            try:
+                self.default_exception_handler(context)
+            except (SystemExit, KeyboardInterrupt):
+                raise
+            except BaseException:
+                # Second protection layer for unexpected errors
+                # in the default implementation, as well as for subclassed
+                # event loops with overloaded "default_exception_handler".
+                print("Exception in default exception handler", file=sys.stderr)
+                traceback.print_exc()
+        else:
+            try:
+                self._exception_handler(self, context)
+            except (SystemExit, KeyboardInterrupt):
+                raise
+            except BaseException as exc:
+                # Exception in the user set custom exception handler.
+                try:
+                    # Let's try default handler.
+                    self.default_exception_handler(
+                        {
+                            "message": "Unhandled error in exception handler",
+                            "exception": exc,
+                            "context": context,
+                        }
+                    )
+                except (SystemExit, KeyboardInterrupt):
+                    raise
+                except BaseException:
+                    # Guard 'default_exception_handler' in case it is
+                    # overloaded.
+                    print(
+                        "Exception in default exception handler "
+                        "while handling an unexpected error "
+                        "in custom exception handler",
+                        file=sys.stderr,
+                    )
+                    traceback.print_exc()
+
 
 class WebLoopPolicy(asyncio.DefaultEventLoopPolicy):  # type: ignore
     """

src/tests/test_webloop.py

@@ -160,3 +160,34 @@ def test_run_in_executor(selenium):
         `);
         """
     )
+
+
+def test_webloop_exception_handler(selenium):
+    selenium.run(
+        """
+        import asyncio
+        async def test():
+            raise Exception("test")
+        asyncio.ensure_future(test())
+        pass
+        """
+    )
+    assert "Task exception was never retrieved" in selenium.logs
+    try:
+        selenium.run(
+            """
+            import asyncio
+            loop = asyncio.get_event_loop()
+            def exception_handler(loop, context):
+                global exc
+                exc = context
+            loop.set_exception_handler(exception_handler)
+
+            async def test():
+                raise Exception("blah")
+            asyncio.ensure_future(test());
+            """
+        )
+        assert selenium.run('exc["exception"].args[0] == "blah"')
+    finally:
+        selenium.run("loop.set_exception_handler(None)")