From 36a7adb5eca9baa56f5ab1af15acac371e20011a Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@python.org>
Date: Sun, 11 Sep 2022 12:26:41 -0700
Subject: [PATCH] GH-96704: Add task.get_context(), use it in
 call_exception_handler()

---
 Lib/asyncio/base_events.py        | 11 ++++++++++-
 Lib/asyncio/tasks.py              |  3 +++
 Modules/_asynciomodule.c          | 13 +++++++++++++
 Modules/clinic/_asynciomodule.c.h | 19 ++++++++++++++++++-
 4 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py
index a675fff1688d27..f8388a476d2b76 100644
--- a/Lib/asyncio/base_events.py
+++ b/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:
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 56a355cbdc70c8..683547f1d217f7 100644
--- a/Lib/asyncio/tasks.py
+++ b/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
 
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index 9d2f83bf6c73c1..0b3b5c59f6c165 100644
--- a/Modules/_asynciomodule.c
+++ b/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 */
 };
diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h
index daf524c3456ca8..ddec54c8d7c2bc 100644
--- a/Modules/clinic/_asynciomodule.c.h
+++ b/Modules/clinic/_asynciomodule.c.h
@@ -772,6 +772,23 @@ _asyncio_Task_get_coro(TaskObj *self, PyObject *Py_UNUSED(ignored))
     return _asyncio_Task_get_coro_impl(self);
 }
 
+PyDoc_STRVAR(_asyncio_Task_get_context__doc__,
+"get_context($self, /)\n"
+"--\n"
+"\n");
+
+#define _ASYNCIO_TASK_GET_CONTEXT_METHODDEF    \
+    {"get_context", (PyCFunction)_asyncio_Task_get_context, METH_NOARGS, _asyncio_Task_get_context__doc__},
+
+static PyObject *
+_asyncio_Task_get_context_impl(TaskObj *self);
+
+static PyObject *
+_asyncio_Task_get_context(TaskObj *self, PyObject *Py_UNUSED(ignored))
+{
+    return _asyncio_Task_get_context_impl(self);
+}
+
 PyDoc_STRVAR(_asyncio_Task_get_name__doc__,
 "get_name($self, /)\n"
 "--\n"
@@ -1172,4 +1189,4 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=459a7c7f21bbc290 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=f117b2246eaf7a55 input=a9049054013a1b77]*/