bpo-47167: Allow overriding a future compliance check in asyncio.Task (GH-32197)

This commit is contained in:
Andrew Svetlov 2022-04-01 04:25:15 +03:00 • committed by GitHub
parent ab89ccff3c
commit d4bb38f82b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 133 additions and 15 deletions

View file

@ -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;

View file

@ -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]*/