pythonGH-96704: Add task.get_context(), use it in call_exception_handler()

Author: gvanrossum

Lib/asyncio/base_events.py

@@ -1794,7 +1794,16 @@ def call_exception_handler(self, context):
                              exc_info=True)
         else:
             try:
-                self._exception_handler(self, context)
+                ctx = None
+                task = context.get("task")
+                if task is None:
+                    task = context.get("future")
+                if task is not None and hasattr(task, "get_context"):
+                    ctx = task.get_context()
+                if ctx is not None and hasattr(ctx, "run"):
+                    ctx.run(self._exception_handler, self, context)
+                else:
+                    self._exception_handler(self, context)
             except (SystemExit, KeyboardInterrupt):
                 raise
             except BaseException as exc:

Lib/asyncio/tasks.py

@@ -139,6 +139,9 @@ def __repr__(self):
     def get_coro(self):
         return self._coro
 
+    def get_context(self):
+        return self._context
+
     def get_name(self):
         return self._name
 

Modules/_asynciomodule.c

@@ -2396,6 +2396,18 @@ _asyncio_Task_get_coro_impl(TaskObj *self)
     return self->task_coro;
 }
 
+/*[clinic input]
+_asyncio.Task.get_context
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio_Task_get_context_impl(TaskObj *self)
+/*[clinic end generated code: output=6996f53d3dc01aef input=87c0b209b8fceeeb]*/
+{
+    Py_INCREF(self->task_context);
+    return self->task_context;
+}
+
 /*[clinic input]
 _asyncio.Task.get_name
 [clinic start generated code]*/
@@ -2523,6 +2535,7 @@ static PyMethodDef TaskType_methods[] = {
     _ASYNCIO_TASK_GET_NAME_METHODDEF
     _ASYNCIO_TASK_SET_NAME_METHODDEF
     _ASYNCIO_TASK_GET_CORO_METHODDEF
+    _ASYNCIO_TASK_GET_CONTEXT_METHODDEF
     {"__class_getitem__", Py_GenericAlias, METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
     {NULL, NULL}        /* Sentinel */
 };