Skip to content

Commit d4bb38f

Browse files
authored
bpo-47167: Allow overriding a future compliance check in asyncio.Task (GH-32197)
1 parent ab89ccf commit d4bb38f

File tree

6 files changed

+133
-15
lines changed

6 files changed

+133
-15
lines changed

Doc/library/asyncio-extending.rst

+16-5
Original file line numberDiff line numberDiff line change
@@ -48,16 +48,27 @@ For this purpose the following, *private* constructors are listed:
4848

4949
.. method:: Future.__init__(*, loop=None)
5050

51-
Create a built-in future instance.
51+
Create a built-in future instance.
5252

53-
*loop* is an optional event loop instance.
53+
*loop* is an optional event loop instance.
5454

5555
.. method:: Task.__init__(coro, *, loop=None, name=None, context=None)
5656

57-
Create a built-in task instance.
57+
Create a built-in task instance.
5858

59-
*loop* is an optional event loop instance. The rest of arguments are described in
60-
:meth:`loop.create_task` description.
59+
*loop* is an optional event loop instance. The rest of arguments are described in
60+
:meth:`loop.create_task` description.
61+
62+
.. versionchanged:: 3.11
63+
64+
*context* argument is added.
65+
66+
.. method:: Tasl._check_future(future)
67+
68+
Return ``True`` if *future* is attached to the same loop as the task, ``False``
69+
otherwise.
70+
71+
.. versionadded:: 3.11
6172

6273

6374
Task lifetime support

Lib/asyncio/tasks.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,10 @@ def uncancel(self):
252252
self._num_cancels_requested -= 1
253253
return self._num_cancels_requested
254254

255+
def _check_future(self, future):
256+
"""Return False if task and future loops are not compatible."""
257+
return futures._get_loop(future) is self._loop
258+
255259
def __step(self, exc=None):
256260
if self.done():
257261
raise exceptions.InvalidStateError(
@@ -292,7 +296,7 @@ def __step(self, exc=None):
292296
blocking = getattr(result, '_asyncio_future_blocking', None)
293297
if blocking is not None:
294298
# Yielded Future must come from Future.__iter__().
295-
if futures._get_loop(result) is not self._loop:
299+
if not self._check_future(result):
296300
new_exc = RuntimeError(
297301
f'Task {self!r} got Future '
298302
f'{result!r} attached to a different loop')

Lib/test/test_asyncio/test_tasks.py

+9-1
Original file line numberDiff line numberDiff line change
@@ -2383,7 +2383,13 @@ def add_done_callback(self, *args, **kwargs):
23832383
return super().add_done_callback(*args, **kwargs)
23842384

23852385
class Task(CommonFuture, BaseTask):
2386-
pass
2386+
def __init__(self, *args, **kwargs):
2387+
self._check_future_called = 0
2388+
super().__init__(*args, **kwargs)
2389+
2390+
def _check_future(self, future):
2391+
self._check_future_called += 1
2392+
return super()._check_future(future)
23872393

23882394
class Future(CommonFuture, BaseFuture):
23892395
pass
@@ -2409,6 +2415,8 @@ async def func():
24092415
dict(fut.calls),
24102416
{'add_done_callback': 1})
24112417

2418+
self.assertEqual(1, task._check_future_called)
2419+
24122420
# Add patched Task & Future back to the test case
24132421
cls.Task = Task
24142422
cls.Future = Future
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Allow overriding a future compliance check in :class:`asyncio.Task`.

Modules/_asynciomodule.c

+64-7
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ _Py_IDENTIFIER(call_soon);
2323
_Py_IDENTIFIER(cancel);
2424
_Py_IDENTIFIER(get_event_loop);
2525
_Py_IDENTIFIER(throw);
26+
_Py_IDENTIFIER(_check_future);
2627

2728

2829
/* State of the _asyncio module */
@@ -1795,6 +1796,8 @@ class _asyncio.Task "TaskObj *" "&Task_Type"
17951796
static int task_call_step_soon(TaskObj *, PyObject *);
17961797
static PyObject * task_wakeup(TaskObj *, PyObject *);
17971798
static PyObject * task_step(TaskObj *, PyObject *);
1799+
static int task_check_future(TaskObj *, PyObject *);
1800+
static int task_check_future_exact(TaskObj *, PyObject *);
17981801

17991802
/* ----- Task._step wrapper */
18001803

@@ -2269,14 +2272,28 @@ Returns the remaining number of cancellation requests.
22692272
static PyObject *
22702273
_asyncio_Task_uncancel_impl(TaskObj *self)
22712274
/*[clinic end generated code: output=58184d236a817d3c input=68f81a4b90b46be2]*/
2272-
/*[clinic end generated code]*/
22732275
{
22742276
if (self->task_num_cancels_requested > 0) {
22752277
self->task_num_cancels_requested -= 1;
22762278
}
22772279
return PyLong_FromLong(self->task_num_cancels_requested);
22782280
}
22792281

2282+
/*[clinic input]
2283+
_asyncio.Task._check_future -> bool
2284+
2285+
future: object
2286+
2287+
Return False if task and future loops are not compatible.
2288+
[clinic start generated code]*/
2289+
2290+
static int
2291+
_asyncio_Task__check_future_impl(TaskObj *self, PyObject *future)
2292+
/*[clinic end generated code: output=a3bfba79295c8d57 input=3b1d6dfd6fe90aa5]*/
2293+
{
2294+
return task_check_future_exact(self, future);
2295+
}
2296+
22802297
/*[clinic input]
22812298
_asyncio.Task.get_stack
22822299
@@ -2502,6 +2519,7 @@ static PyMethodDef TaskType_methods[] = {
25022519
_ASYNCIO_TASK_CANCEL_METHODDEF
25032520
_ASYNCIO_TASK_CANCELLING_METHODDEF
25042521
_ASYNCIO_TASK_UNCANCEL_METHODDEF
2522+
_ASYNCIO_TASK__CHECK_FUTURE_METHODDEF
25052523
_ASYNCIO_TASK_GET_STACK_METHODDEF
25062524
_ASYNCIO_TASK_PRINT_STACK_METHODDEF
25072525
_ASYNCIO_TASK__MAKE_CANCELLED_ERROR_METHODDEF
@@ -2569,6 +2587,43 @@ TaskObj_dealloc(PyObject *self)
25692587
Py_TYPE(task)->tp_free(task);
25702588
}
25712589

2590+
static int
2591+
task_check_future_exact(TaskObj *task, PyObject *future)
2592+
{
2593+
int res;
2594+
if (Future_CheckExact(future) || Task_CheckExact(future)) {
2595+
FutureObj *fut = (FutureObj *)future;
2596+
res = (fut->fut_loop == task->task_loop);
2597+
} else {
2598+
PyObject *oloop = get_future_loop(future);
2599+
if (oloop == NULL) {
2600+
return -1;
2601+
}
2602+
res = (oloop == task->task_loop);
2603+
Py_DECREF(oloop);
2604+
}
2605+
return res;
2606+
}
2607+
2608+
2609+
static int
2610+
task_check_future(TaskObj *task, PyObject *future)
2611+
{
2612+
if (Task_CheckExact(task)) {
2613+
return task_check_future_exact(task, future);
2614+
} else {
2615+
PyObject * ret = _PyObject_CallMethodIdOneArg((PyObject *)task,
2616+
&PyId__check_future,
2617+
future);
2618+
if (ret == NULL) {
2619+
return -1;
2620+
}
2621+
int is_true = PyObject_IsTrue(ret);
2622+
Py_DECREF(ret);
2623+
return is_true;
2624+
}
2625+
}
2626+
25722627
static int
25732628
task_call_step_soon(TaskObj *task, PyObject *arg)
25742629
{
@@ -2790,7 +2845,11 @@ task_step_impl(TaskObj *task, PyObject *exc)
27902845
FutureObj *fut = (FutureObj*)result;
27912846

27922847
/* Check if `result` future is attached to a different loop */
2793-
if (fut->fut_loop != task->task_loop) {
2848+
res = task_check_future(task, result);
2849+
if (res == -1) {
2850+
goto fail;
2851+
}
2852+
if (res == 0) {
27942853
goto different_loop;
27952854
}
27962855

@@ -2862,15 +2921,13 @@ task_step_impl(TaskObj *task, PyObject *exc)
28622921
}
28632922

28642923
/* Check if `result` future is attached to a different loop */
2865-
PyObject *oloop = get_future_loop(result);
2866-
if (oloop == NULL) {
2924+
res = task_check_future(task, result);
2925+
if (res == -1) {
28672926
goto fail;
28682927
}
2869-
if (oloop != task->task_loop) {
2870-
Py_DECREF(oloop);
2928+
if (res == 0) {
28712929
goto different_loop;
28722930
}
2873-
Py_DECREF(oloop);
28742931

28752932
if (!blocking) {
28762933
goto yield_insteadof_yf;

Modules/clinic/_asynciomodule.c.h

+38-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)