https://github.com/python/cpython/commit/d4bb38f82bf18b00db3129031ce4969b6f0... commit: d4bb38f82bf18b00db3129031ce4969b6f0caab9 branch: main author: Andrew Svetlov <andrew.svetlov@gmail.com> committer: asvetlov <andrew.svetlov@gmail.com> date: 2022-04-01T04:25:15+03:00 summary: bpo-47167: Allow overriding a future compliance check in asyncio.Task (GH-32197) files: A Misc/NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst M Doc/library/asyncio-extending.rst M Lib/asyncio/tasks.py M Lib/test/test_asyncio/test_tasks.py M Modules/_asynciomodule.c M Modules/clinic/_asynciomodule.c.h diff --git a/Doc/library/asyncio-extending.rst b/Doc/library/asyncio-extending.rst index 619723e61b5f9..215d215bb14fe 100644 --- a/Doc/library/asyncio-extending.rst +++ b/Doc/library/asyncio-extending.rst @@ -48,16 +48,27 @@ For this purpose the following, *private* constructors are listed: .. method:: Future.__init__(*, loop=None) -Create a built-in future instance. + Create a built-in future instance. -*loop* is an optional event loop instance. + *loop* is an optional event loop instance. .. method:: Task.__init__(coro, *, loop=None, name=None, context=None) -Create a built-in task instance. + Create a built-in task instance. -*loop* is an optional event loop instance. The rest of arguments are described in -:meth:`loop.create_task` description. + *loop* is an optional event loop instance. The rest of arguments are described in + :meth:`loop.create_task` description. + + .. versionchanged:: 3.11 + + *context* argument is added. + +.. method:: Tasl._check_future(future) + + Return ``True`` if *future* is attached to the same loop as the task, ``False`` + otherwise. + + .. versionadded:: 3.11 Task lifetime support diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 27fe58da15136..3952b5f2a7743 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -252,6 +252,10 @@ def uncancel(self): self._num_cancels_requested -= 1 return self._num_cancels_requested + def _check_future(self, future): + """Return False if task and future loops are not compatible.""" + return futures._get_loop(future) is self._loop + def __step(self, exc=None): if self.done(): raise exceptions.InvalidStateError( @@ -292,7 +296,7 @@ def __step(self, exc=None): blocking = getattr(result, '_asyncio_future_blocking', None) if blocking is not None: # Yielded Future must come from Future.__iter__(). - if futures._get_loop(result) is not self._loop: + if not self._check_future(result): new_exc = RuntimeError( f'Task {self!r} got Future ' f'{result!r} attached to a different loop') diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 8df1957bbe9e7..80afb27351362 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -2383,7 +2383,13 @@ def add_done_callback(self, *args, **kwargs): return super().add_done_callback(*args, **kwargs) class Task(CommonFuture, BaseTask): - pass + def __init__(self, *args, **kwargs): + self._check_future_called = 0 + super().__init__(*args, **kwargs) + + def _check_future(self, future): + self._check_future_called += 1 + return super()._check_future(future) class Future(CommonFuture, BaseFuture): pass @@ -2409,6 +2415,8 @@ async def func(): dict(fut.calls), {'add_done_callback': 1}) + self.assertEqual(1, task._check_future_called) + # Add patched Task & Future back to the test case cls.Task = Task cls.Future = Future diff --git a/Misc/NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst b/Misc/NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst new file mode 100644 index 0000000000000..a37dd25810f2e --- /dev/null +++ b/Misc/NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst @@ -0,0 +1 @@ +Allow overriding a future compliance check in :class:`asyncio.Task`. diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c index 632a4465c224a..d8d3da91cdd8e 100644 --- a/Modules/_asynciomodule.c +++ b/Modules/_asynciomodule.c @@ -23,6 +23,7 @@ _Py_IDENTIFIER(call_soon); _Py_IDENTIFIER(cancel); _Py_IDENTIFIER(get_event_loop); _Py_IDENTIFIER(throw); +_Py_IDENTIFIER(_check_future); /* State of the _asyncio module */ @@ -1795,6 +1796,8 @@ class _asyncio.Task "TaskObj *" "&Task_Type" static int task_call_step_soon(TaskObj *, PyObject *); static PyObject * task_wakeup(TaskObj *, PyObject *); static PyObject * task_step(TaskObj *, PyObject *); +static int task_check_future(TaskObj *, PyObject *); +static int task_check_future_exact(TaskObj *, PyObject *); /* ----- Task._step wrapper */ @@ -2269,7 +2272,6 @@ Returns the remaining number of cancellation requests. static PyObject * _asyncio_Task_uncancel_impl(TaskObj *self) /*[clinic end generated code: output=58184d236a817d3c input=68f81a4b90b46be2]*/ -/*[clinic end generated code]*/ { if (self->task_num_cancels_requested > 0) { self->task_num_cancels_requested -= 1; @@ -2277,6 +2279,21 @@ _asyncio_Task_uncancel_impl(TaskObj *self) return PyLong_FromLong(self->task_num_cancels_requested); } +/*[clinic input] +_asyncio.Task._check_future -> bool + + future: object + +Return False if task and future loops are not compatible. +[clinic start generated code]*/ + +static int +_asyncio_Task__check_future_impl(TaskObj *self, PyObject *future) +/*[clinic end generated code: output=a3bfba79295c8d57 input=3b1d6dfd6fe90aa5]*/ +{ + return task_check_future_exact(self, future); +} + /*[clinic input] _asyncio.Task.get_stack @@ -2502,6 +2519,7 @@ static PyMethodDef TaskType_methods[] = { _ASYNCIO_TASK_CANCEL_METHODDEF _ASYNCIO_TASK_CANCELLING_METHODDEF _ASYNCIO_TASK_UNCANCEL_METHODDEF + _ASYNCIO_TASK__CHECK_FUTURE_METHODDEF _ASYNCIO_TASK_GET_STACK_METHODDEF _ASYNCIO_TASK_PRINT_STACK_METHODDEF _ASYNCIO_TASK__MAKE_CANCELLED_ERROR_METHODDEF @@ -2569,6 +2587,43 @@ TaskObj_dealloc(PyObject *self) Py_TYPE(task)->tp_free(task); } +static int +task_check_future_exact(TaskObj *task, PyObject *future) +{ + int res; + if (Future_CheckExact(future) || Task_CheckExact(future)) { + FutureObj *fut = (FutureObj *)future; + res = (fut->fut_loop == task->task_loop); + } else { + PyObject *oloop = get_future_loop(future); + if (oloop == NULL) { + return -1; + } + res = (oloop == task->task_loop); + Py_DECREF(oloop); + } + return res; +} + + +static int +task_check_future(TaskObj *task, PyObject *future) +{ + if (Task_CheckExact(task)) { + return task_check_future_exact(task, future); + } else { + PyObject * ret = _PyObject_CallMethodIdOneArg((PyObject *)task, + &PyId__check_future, + future); + if (ret == NULL) { + return -1; + } + int is_true = PyObject_IsTrue(ret); + Py_DECREF(ret); + return is_true; + } +} + static int task_call_step_soon(TaskObj *task, PyObject *arg) { @@ -2790,7 +2845,11 @@ task_step_impl(TaskObj *task, PyObject *exc) FutureObj *fut = (FutureObj*)result; /* Check if `result` future is attached to a different loop */ - if (fut->fut_loop != task->task_loop) { + res = task_check_future(task, result); + if (res == -1) { + goto fail; + } + if (res == 0) { goto different_loop; } @@ -2862,15 +2921,13 @@ task_step_impl(TaskObj *task, PyObject *exc) } /* Check if `result` future is attached to a different loop */ - PyObject *oloop = get_future_loop(result); - if (oloop == NULL) { + res = task_check_future(task, result); + if (res == -1) { goto fail; } - if (oloop != task->task_loop) { - Py_DECREF(oloop); + if (res == 0) { goto different_loop; } - Py_DECREF(oloop); if (!blocking) { goto yield_insteadof_yf; diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h index 4b64367a3f631..163b0f95691b3 100644 --- a/Modules/clinic/_asynciomodule.c.h +++ b/Modules/clinic/_asynciomodule.c.h @@ -466,6 +466,43 @@ _asyncio_Task_uncancel(TaskObj *self, PyObject *Py_UNUSED(ignored)) return _asyncio_Task_uncancel_impl(self); } +PyDoc_STRVAR(_asyncio_Task__check_future__doc__, +"_check_future($self, /, future)\n" +"--\n" +"\n" +"Return False if task and future loops are not compatible."); + +#define _ASYNCIO_TASK__CHECK_FUTURE_METHODDEF \ + {"_check_future", (PyCFunction)(void(*)(void))_asyncio_Task__check_future, METH_FASTCALL|METH_KEYWORDS, _asyncio_Task__check_future__doc__}, + +static int +_asyncio_Task__check_future_impl(TaskObj *self, PyObject *future); + +static PyObject * +_asyncio_Task__check_future(TaskObj *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) +{ + PyObject *return_value = NULL; + static const char * const _keywords[] = {"future", NULL}; + static _PyArg_Parser _parser = {NULL, _keywords, "_check_future", 0}; + PyObject *argsbuf[1]; + PyObject *future; + int _return_value; + + args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 1, 1, 0, argsbuf); + if (!args) { + goto exit; + } + future = args[0]; + _return_value = _asyncio_Task__check_future_impl(self, future); + if ((_return_value == -1) && PyErr_Occurred()) { + goto exit; + } + return_value = PyBool_FromLong((long)_return_value); + +exit: + return return_value; +} + PyDoc_STRVAR(_asyncio_Task_get_stack__doc__, "get_stack($self, /, *, limit=None)\n" "--\n" @@ -890,4 +927,4 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs, exit: return return_value; } -/*[clinic end generated code: output=64b3836574e8a18c input=a9049054013a1b77]*/ +/*[clinic end generated code: output=fdb7129263a8712e input=a9049054013a1b77]*/