Expose Task._step

Summary:
why is this needed? see https://discuss.python.org/t/request-can-we-get-a-c-api-hook-into-pycontext-enter-and-pycontext-exit/51730

based on D38173091 and D46824875

Reviewed By: fried

Differential Revision: D56453986

fbshipit-source-id: 1a965de8369ff3be70b5fe879d41237ecdddc730

Include/internal/pycore_global_strings.h

@@ -262,6 +262,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(_showwarnmsg)
         STRUCT_FOR_ID(_shutdown)
         STRUCT_FOR_ID(_slotnames)
+        STRUCT_FOR_ID(_step)
         STRUCT_FOR_ID(_strptime_datetime)
         STRUCT_FOR_ID(_swappedbytes_)
         STRUCT_FOR_ID(_type_)
@@ -402,6 +403,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(errors)
         STRUCT_FOR_ID(event)
         STRUCT_FOR_ID(eventmask)
+        STRUCT_FOR_ID(exc)
         STRUCT_FOR_ID(exc_type)
         STRUCT_FOR_ID(exc_value)
         STRUCT_FOR_ID(excepthook)

Lib/asyncio/tasks.py

@@ -380,6 +380,11 @@ def __step_run_and_handle_result(self, exc):
         finally:
             self = None  # Needed to break cycles when an exception occurs.
 
+    # START META PATCH (Task._step override)
+    # Needed to be compatible with the C version
+    _step = __step
+    # END META PATCH
+
     def __wakeup(self, future):
         try:
             future.result()

Lib/test/test_asyncio/test_tasks.py

@@ -3365,6 +3365,75 @@ async def coro():
         self.assertEqual(result, 11)
 
 
+# START META PATCH (Task._step override tests)
+class MetaAsyncio(test_utils.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.loop = asyncio.new_event_loop()
+        self.set_event_loop(self.loop)
+
+    def tearDown(self):
+        self.loop.close()
+        self.loop = None
+        super().tearDown()
+
+    def test_override_step_no_exc(self):
+        result = ""
+
+        async def f():
+            nonlocal result
+            result += "2"
+
+        class CustomTask(asyncio.Task):
+            def _step(self, exn=None):
+                nonlocal result
+                result += "1"
+                super()._step(exn)
+                result += "3"
+
+        self.loop.run_until_complete(CustomTask(f(), loop=self.loop))
+        self.assertEqual(result, "123")
+
+    def test_override_step_has_exc(self):
+        result = ""
+
+        # we want 'await fut' to not be finished synchronously
+        # in order to enter '_step' for the second time with exception.
+        # If we just do 'set_exception' and then await - it will be finished
+        # eagerly since future is fulfiled at th suspension point.
+        # What we want is to set error after the point when asyncio infrastructure
+        # will decide to suspend execution. Simplest way to do it is to  set error at the point
+        # when done callback is set.
+        class F(asyncio.Future):
+            def add_done_callback(self, callback, *, context=None):
+                super().add_done_callback(callback, context=context)
+                self.set_exception(ValueError("ERROR!"))
+
+        fut = F(loop=self.loop)
+
+        async def f():
+            nonlocal result
+            result += "|f"
+            result += await fut
+
+        class CustomTask(asyncio.Task):
+            def _step(self, exn=None):
+                nonlocal result
+                result += "|before"
+                result += "|" + str(exn)
+                super()._step(exn)
+                result += "|after"
+
+        try:
+            self.loop.run_until_complete(CustomTask(f(), loop=self.loop))
+        except ValueError as e:
+            self.assertEqual(e.args[0], "ERROR!")
+        else:
+            self.fail("ValueError expected")
+        self.assertEqual(result, "|before|None|f|after|before|ERROR!|after")
+# END META PATCH
+
+
 class CompatibilityTests(test_utils.TestCase):
     # Tests for checking a bridge between old-styled coroutines
     # and async/await syntax

Modules/_asynciomodule.c

@@ -1856,6 +1856,21 @@ TaskStepMethWrapper_dealloc(TaskStepMethWrapper *o)
     Py_DECREF(tp);
 }
 
+// START META PATCH (Task._step override)
+static inline PyObject *
+task_call_step(asyncio_state *state, TaskObj *task, PyObject *arg)
+{
+    if (Task_CheckExact(state, task)) {
+        return task_step(state, task, arg);
+    }
+    else {
+        // `task` is a subclass of asyncio.Task
+        return PyObject_CallMethodObjArgs(
+            (PyObject*)task, &_Py_ID(_step), arg, NULL);
+    }
+}
+// END META PATCH
+
 static PyObject *
 TaskStepMethWrapper_call(TaskStepMethWrapper *o,
                          PyObject *args, PyObject *kwds)
@@ -1869,7 +1884,9 @@ TaskStepMethWrapper_call(TaskStepMethWrapper *o,
         return NULL;
     }
     asyncio_state *state = get_asyncio_state_by_def((PyObject *)o);
-    return task_step(state, o->sw_task, o->sw_arg);
+    // START META PATCH (Task._step override)
+    return task_call_step(state, o->sw_task, o->sw_arg);
+    // END META PATCH
 }
 
 static int
@@ -2487,6 +2504,22 @@ _asyncio_Task_set_exception(TaskObj *self, PyObject *exception)
     return NULL;
 }
 
+// START META PATCH (Task._step override)
+/*[clinic input]
+_asyncio.Task._step
+
+    exc: object = None
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio_Task__step_impl(TaskObj *self, PyObject *exc)
+/*[clinic end generated code: output=7ed23f0cefd5ae42 input=1e19a985ace87ca4]*/
+{
+    asyncio_state *state = get_asyncio_state_by_def((PyObject *)self);
+    return task_step(state, self, exc == Py_None ? NULL : exc);
+}
+// END META PATCH
+
 /*[clinic input]
 _asyncio.Task.get_coro
 [clinic start generated code]*/
@@ -2637,6 +2670,9 @@ static PyMethodDef TaskType_methods[] = {
     _ASYNCIO_TASK_SET_NAME_METHODDEF
     _ASYNCIO_TASK_GET_CORO_METHODDEF
     _ASYNCIO_TASK_GET_CONTEXT_METHODDEF
+    // START META PATCH (Task._step override)
+    _ASYNCIO_TASK__STEP_METHODDEF
+    // END META PATCH
     {"__class_getitem__", Py_GenericAlias, METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
     {NULL, NULL}        /* Sentinel */
 };
@@ -3219,10 +3255,14 @@ task_wakeup(TaskObj *task, PyObject *o)
             break; /* exception raised */
         case 0:
             Py_DECREF(fut_result);
-            return task_step(state, task, NULL);
+            // START META PATCH (Task._step override)
+            return task_call_step(state, task, NULL);
+            // END META PATCH
         default:
             assert(res == 1);
-            result = task_step(state, task, fut_result);
+            // START META PATCH (Task._step override)
+            result = task_call_step(state, task, fut_result);
+            // END META PATCH
             Py_DECREF(fut_result);
             return result;
         }
@@ -3231,15 +3271,19 @@ task_wakeup(TaskObj *task, PyObject *o)
         PyObject *fut_result = PyObject_CallMethod(o, "result", NULL);
         if (fut_result != NULL) {
             Py_DECREF(fut_result);
-            return task_step(state, task, NULL);
+            // START META PATCH (Task._step override)
+            return task_call_step(state, task, NULL);
+            // END META PATCH
         }
         /* exception raised */
     }
 
     PyObject *exc = PyErr_GetRaisedException();
     assert(exc);
 
-    result = task_step(state, task, exc);
+    // START META PATCH (Task._step override)
+    result = task_call_step(state, task, exc);
+    // END META PATCH
 
     Py_DECREF(exc);