diff --git a/Include/internal/pycore_global_objects_fini_generated.h b/Include/internal/pycore_global_objects_fini_generated.h index decdba6c0d7..efb52d7b51f 100644 --- a/Include/internal/pycore_global_objects_fini_generated.h +++ b/Include/internal/pycore_global_objects_fini_generated.h @@ -773,6 +773,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) { _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(_showwarnmsg)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(_shutdown)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(_slotnames)); + _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(_step)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(_strptime_datetime)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(_swappedbytes_)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(_type_)); @@ -913,6 +914,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) { _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(errors)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(event)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(eventmask)); + _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(exc)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(exc_type)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(exc_value)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(excepthook)); diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h index c5f7532effd..23af9dcac44 100644 --- a/Include/internal/pycore_global_strings.h +++ b/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) diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h index fc1c1d29f42..8bee4029d5d 100644 --- a/Include/internal/pycore_runtime_init_generated.h +++ b/Include/internal/pycore_runtime_init_generated.h @@ -768,6 +768,7 @@ extern "C" { INIT_ID(_showwarnmsg), \ INIT_ID(_shutdown), \ INIT_ID(_slotnames), \ + INIT_ID(_step), \ INIT_ID(_strptime_datetime), \ INIT_ID(_swappedbytes_), \ INIT_ID(_type_), \ @@ -908,6 +909,7 @@ extern "C" { INIT_ID(errors), \ INIT_ID(event), \ INIT_ID(eventmask), \ + INIT_ID(exc), \ INIT_ID(exc_type), \ INIT_ID(exc_value), \ INIT_ID(excepthook), \ diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h index b3c560290bf..fb467ca29a1 100644 --- a/Include/internal/pycore_unicodeobject_generated.h +++ b/Include/internal/pycore_unicodeobject_generated.h @@ -627,6 +627,9 @@ _PyUnicode_InitStaticStrings(PyInterpreterState *interp) { string = &_Py_ID(_slotnames); assert(_PyUnicode_CheckConsistency(string, 1)); _PyUnicode_InternInPlace(interp, &string); + string = &_Py_ID(_step); + assert(_PyUnicode_CheckConsistency(string, 1)); + _PyUnicode_InternInPlace(interp, &string); string = &_Py_ID(_strptime_datetime); assert(_PyUnicode_CheckConsistency(string, 1)); _PyUnicode_InternInPlace(interp, &string); @@ -1047,6 +1050,9 @@ _PyUnicode_InitStaticStrings(PyInterpreterState *interp) { string = &_Py_ID(eventmask); assert(_PyUnicode_CheckConsistency(string, 1)); _PyUnicode_InternInPlace(interp, &string); + string = &_Py_ID(exc); + assert(_PyUnicode_CheckConsistency(string, 1)); + _PyUnicode_InternInPlace(interp, &string); string = &_Py_ID(exc_type); assert(_PyUnicode_CheckConsistency(string, 1)); _PyUnicode_InternInPlace(interp, &string); diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 0b22e28d8e0..881637f660e 100644 --- a/Lib/asyncio/tasks.py +++ b/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() diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 6e8a51ce255..72dc2195547 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/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 diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c index a465090bfaa..fc522dab122 100644 --- a/Modules/_asynciomodule.c +++ b/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,7 +3271,9 @@ 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 */ } @@ -3239,7 +3281,9 @@ task_wakeup(TaskObj *task, PyObject *o) 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); diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h index 860d55cb3bb..0a06a035182 100644 --- a/Modules/clinic/_asynciomodule.c.h +++ b/Modules/clinic/_asynciomodule.c.h @@ -881,6 +881,65 @@ PyDoc_STRVAR(_asyncio_Task_set_exception__doc__, #define _ASYNCIO_TASK_SET_EXCEPTION_METHODDEF \ {"set_exception", (PyCFunction)_asyncio_Task_set_exception, METH_O, _asyncio_Task_set_exception__doc__}, +PyDoc_STRVAR(_asyncio_Task__step__doc__, +"_step($self, /, exc=None)\n" +"--\n" +"\n"); + +#define _ASYNCIO_TASK__STEP_METHODDEF \ + {"_step", _PyCFunction_CAST(_asyncio_Task__step), METH_FASTCALL|METH_KEYWORDS, _asyncio_Task__step__doc__}, + +static PyObject * +_asyncio_Task__step_impl(TaskObj *self, PyObject *exc); + +static PyObject * +_asyncio_Task__step(TaskObj *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) +{ + PyObject *return_value = NULL; + #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE) + + #define NUM_KEYWORDS 1 + static struct { + PyGC_Head _this_is_not_used; + PyObject_VAR_HEAD + PyObject *ob_item[NUM_KEYWORDS]; + } _kwtuple = { + .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS) + .ob_item = { &_Py_ID(exc), }, + }; + #undef NUM_KEYWORDS + #define KWTUPLE (&_kwtuple.ob_base.ob_base) + + #else // !Py_BUILD_CORE + # define KWTUPLE NULL + #endif // !Py_BUILD_CORE + + static const char * const _keywords[] = {"exc", NULL}; + static _PyArg_Parser _parser = { + .keywords = _keywords, + .fname = "_step", + .kwtuple = KWTUPLE, + }; + #undef KWTUPLE + PyObject *argsbuf[1]; + Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0; + PyObject *exc = Py_None; + + args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 0, 1, 0, argsbuf); + if (!args) { + goto exit; + } + if (!noptargs) { + goto skip_optional_pos; + } + exc = args[0]; +skip_optional_pos: + return_value = _asyncio_Task__step_impl(self, exc); + +exit: + return return_value; +} + PyDoc_STRVAR(_asyncio_Task_get_coro__doc__, "get_coro($self, /)\n" "--\n" @@ -1487,4 +1546,4 @@ _asyncio_current_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs, exit: return return_value; } -/*[clinic end generated code: output=127ba6153250d769 input=a9049054013a1b77]*/ +/*[clinic end generated code: output=05c15fa7a7439615 input=a9049054013a1b77]*/