bpo-32436: Implement PEP 567 (#5027)

1 files changed, 154 insertions, 55 deletions

diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index 22ce32c..f77ec99 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -36,6 +36,7 @@ static PyObject *asyncio_task_print_stack_func;
 static PyObject *asyncio_task_repr_info_func;
 static PyObject *asyncio_InvalidStateError;
 static PyObject *asyncio_CancelledError;
+static PyObject *context_kwname;
 
 
 /* WeakSet containing all alive tasks. */
@@ -59,6 +60,7 @@ typedef enum {
     PyObject_HEAD                                                           \
     PyObject *prefix##_loop;                                                \
     PyObject *prefix##_callback0;                                           \
+    PyContext *prefix##_context0;                                           \
     PyObject *prefix##_callbacks;                                           \
     PyObject *prefix##_exception;                                           \
     PyObject *prefix##_result;                                              \
@@ -77,6 +79,7 @@ typedef struct {
     FutureObj_HEAD(task)
     PyObject *task_fut_waiter;
     PyObject *task_coro;
+    PyContext *task_context;
     int task_must_cancel;
     int task_log_destroy_pending;
 } TaskObj;
@@ -336,11 +339,38 @@ get_event_loop(void)
 
 
 static int
-call_soon(PyObject *loop, PyObject *func, PyObject *arg)
+call_soon(PyObject *loop, PyObject *func, PyObject *arg, PyContext *ctx)
 {
     PyObject *handle;
-    handle = _PyObject_CallMethodIdObjArgs(
-        loop, &PyId_call_soon, func, arg, NULL);
+    PyObject *stack[3];
+    Py_ssize_t nargs;
+
+    if (ctx == NULL) {
+        handle = _PyObject_CallMethodIdObjArgs(
+            loop, &PyId_call_soon, func, arg, NULL);
+    }
+    else {
+        /* Use FASTCALL to pass a keyword-only argument to call_soon */
+
+        PyObject *callable = _PyObject_GetAttrId(loop, &PyId_call_soon);
+        if (callable == NULL) {
+            return -1;
+        }
+
+        /* All refs in 'stack' are borrowed. */
+        nargs = 1;
+        stack[0] = func;
+        if (arg != NULL) {
+            stack[1] = arg;
+            nargs++;
+        }
+        stack[nargs] = (PyObject *)ctx;
+
+        handle = _PyObject_FastCallKeywords(
+            callable, stack, nargs, context_kwname);
+        Py_DECREF(callable);
+    }
+
     if (handle == NULL) {
         return -1;
     }
@@ -387,8 +417,11 @@ future_schedule_callbacks(FutureObj *fut)
         /* There's a 1st callback */
 
         int ret = call_soon(
-            fut->fut_loop, fut->fut_callback0, (PyObject *)fut);
+            fut->fut_loop, fut->fut_callback0,
+            (PyObject *)fut, fut->fut_context0);
+
         Py_CLEAR(fut->fut_callback0);
+        Py_CLEAR(fut->fut_context0);
         if (ret) {
             /* If an error occurs in pure-Python implementation,
                all callbacks are cleared. */
@@ -413,9 +446,11 @@ future_schedule_callbacks(FutureObj *fut)
     }
 
     for (i = 0; i < len; i++) {
-        PyObject *cb = PyList_GET_ITEM(fut->fut_callbacks, i);
+        PyObject *cb_tup = PyList_GET_ITEM(fut->fut_callbacks, i);
+        PyObject *cb = PyTuple_GET_ITEM(cb_tup, 0);
+        PyObject *ctx = PyTuple_GET_ITEM(cb_tup, 1);
 
-        if (call_soon(fut->fut_loop, cb, (PyObject *)fut)) {
+        if (call_soon(fut->fut_loop, cb, (PyObject *)fut, (PyContext *)ctx)) {
             /* If an error occurs in pure-Python implementation,
                all callbacks are cleared. */
             Py_CLEAR(fut->fut_callbacks);
@@ -462,6 +497,7 @@ future_init(FutureObj *fut, PyObject *loop)
     }
 
     fut->fut_callback0 = NULL;
+    fut->fut_context0 = NULL;
     fut->fut_callbacks = NULL;
 
     return 0;
@@ -566,7 +602,7 @@ future_get_result(FutureObj *fut, PyObject **result)
 }
 
 static PyObject *
-future_add_done_callback(FutureObj *fut, PyObject *arg)
+future_add_done_callback(FutureObj *fut, PyObject *arg, PyContext *ctx)
 {
     if (!future_is_alive(fut)) {
         PyErr_SetString(PyExc_RuntimeError, "uninitialized Future object");
@@ -576,7 +612,7 @@ future_add_done_callback(FutureObj *fut, PyObject *arg)
     if (fut->fut_state != STATE_PENDING) {
         /* The future is done/cancelled, so schedule the callback
            right away. */
-        if (call_soon(fut->fut_loop, arg, (PyObject*) fut)) {
+        if (call_soon(fut->fut_loop, arg, (PyObject*) fut, ctx)) {
             return NULL;
         }
     }
@@ -602,24 +638,38 @@ future_add_done_callback(FutureObj *fut, PyObject *arg)
                 with a new list and add the new callback to it.
         */
 
-        if (fut->fut_callbacks != NULL) {
-            int err = PyList_Append(fut->fut_callbacks, arg);
-            if (err != 0) {
-                return NULL;
-            }
-        }
-        else if (fut->fut_callback0 == NULL) {
+        if (fut->fut_callbacks == NULL && fut->fut_callback0 == NULL) {
             Py_INCREF(arg);
             fut->fut_callback0 = arg;
+            Py_INCREF(ctx);
+            fut->fut_context0 = ctx;
         }
         else {
-            fut->fut_callbacks = PyList_New(1);
-            if (fut->fut_callbacks == NULL) {
+            PyObject *tup = PyTuple_New(2);
+            if (tup == NULL) {
                 return NULL;
             }
-
             Py_INCREF(arg);
-            PyList_SET_ITEM(fut->fut_callbacks, 0, arg);
+            PyTuple_SET_ITEM(tup, 0, arg);
+            Py_INCREF(ctx);
+            PyTuple_SET_ITEM(tup, 1, (PyObject *)ctx);
+
+            if (fut->fut_callbacks != NULL) {
+                int err = PyList_Append(fut->fut_callbacks, tup);
+                if (err) {
+                    Py_DECREF(tup);
+                    return NULL;
+                }
+                Py_DECREF(tup);
+            }
+            else {
+                fut->fut_callbacks = PyList_New(1);
+                if (fut->fut_callbacks == NULL) {
+                    return NULL;
+                }
+
+                PyList_SET_ITEM(fut->fut_callbacks, 0, tup);  /* borrow */
+            }
         }
     }
 
@@ -676,6 +726,7 @@ FutureObj_clear(FutureObj *fut)
 {
     Py_CLEAR(fut->fut_loop);
     Py_CLEAR(fut->fut_callback0);
+    Py_CLEAR(fut->fut_context0);
     Py_CLEAR(fut->fut_callbacks);
     Py_CLEAR(fut->fut_result);
     Py_CLEAR(fut->fut_exception);
@@ -689,6 +740,7 @@ FutureObj_traverse(FutureObj *fut, visitproc visit, void *arg)
 {
     Py_VISIT(fut->fut_loop);
     Py_VISIT(fut->fut_callback0);
+    Py_VISIT(fut->fut_context0);
     Py_VISIT(fut->fut_callbacks);
     Py_VISIT(fut->fut_result);
     Py_VISIT(fut->fut_exception);
@@ -821,6 +873,8 @@ _asyncio.Future.add_done_callback
 
     fn: object
     /
+    *
+    context: object = NULL
 
 Add a callback to be run when the future becomes done.
 
@@ -830,10 +884,21 @@ scheduled with call_soon.
 [clinic start generated code]*/
 
 static PyObject *
-_asyncio_Future_add_done_callback(FutureObj *self, PyObject *fn)
-/*[clinic end generated code: output=819e09629b2ec2b5 input=8f818b39990b027d]*/
+_asyncio_Future_add_done_callback_impl(FutureObj *self, PyObject *fn,
+                                       PyObject *context)
+/*[clinic end generated code: output=7ce635bbc9554c1e input=15ab0693a96e9533]*/
 {
-    return future_add_done_callback(self, fn);
+    if (context == NULL) {
+        context = (PyObject *)PyContext_CopyCurrent();
+        if (context == NULL) {
+            return NULL;
+        }
+        PyObject *res = future_add_done_callback(
+            self, fn, (PyContext *)context);
+        Py_DECREF(context);
+        return res;
+    }
+    return future_add_done_callback(self, fn, (PyContext *)context);
 }
 
 /*[clinic input]
@@ -865,6 +930,7 @@ _asyncio_Future_remove_done_callback(FutureObj *self, PyObject *fn)
         if (cmp == 1) {
             /* callback0 == fn */
             Py_CLEAR(self->fut_callback0);
+            Py_CLEAR(self->fut_context0);
             cleared_callback0 = 1;
         }
     }
@@ -880,8 +946,9 @@ _asyncio_Future_remove_done_callback(FutureObj *self, PyObject *fn)
     }
 
     if (len == 1) {
+        PyObject *cb_tup = PyList_GET_ITEM(self->fut_callbacks, 0);
         int cmp = PyObject_RichCompareBool(
-            fn, PyList_GET_ITEM(self->fut_callbacks, 0), Py_EQ);
+            fn, PyTuple_GET_ITEM(cb_tup, 0), Py_EQ);
         if (cmp == -1) {
             return NULL;
         }
@@ -903,7 +970,7 @@ _asyncio_Future_remove_done_callback(FutureObj *self, PyObject *fn)
         int ret;
         PyObject *item = PyList_GET_ITEM(self->fut_callbacks, i);
         Py_INCREF(item);
-        ret = PyObject_RichCompareBool(fn, item, Py_EQ);
+        ret = PyObject_RichCompareBool(fn, PyTuple_GET_ITEM(item, 0), Py_EQ);
         if (ret == 0) {
             if (j < len) {
                 PyList_SET_ITEM(newlist, j, item);
@@ -1081,47 +1148,49 @@ static PyObject *
 FutureObj_get_callbacks(FutureObj *fut)
 {
     Py_ssize_t i;
-    Py_ssize_t len;
-    PyObject *new_list;
 
     ENSURE_FUTURE_ALIVE(fut)
 
-    if (fut->fut_callbacks == NULL) {
-        if (fut->fut_callback0 == NULL) {
+    if (fut->fut_callback0 == NULL) {
+        if (fut->fut_callbacks == NULL) {
             Py_RETURN_NONE;
         }
-        else {
-            new_list = PyList_New(1);
-            if (new_list == NULL) {
-                return NULL;
-            }
-            Py_INCREF(fut->fut_callback0);
-            PyList_SET_ITEM(new_list, 0, fut->fut_callback0);
-            return new_list;
-        }
-    }
 
-    assert(fut->fut_callbacks != NULL);
-
-    if (fut->fut_callback0 == NULL) {
         Py_INCREF(fut->fut_callbacks);
         return fut->fut_callbacks;
     }
 
-    assert(fut->fut_callback0 != NULL);
+    Py_ssize_t len = 1;
+    if (fut->fut_callbacks != NULL) {
+        len += PyList_GET_SIZE(fut->fut_callbacks);
+    }
 
-    len = PyList_GET_SIZE(fut->fut_callbacks);
-    new_list = PyList_New(len + 1);
+
+    PyObject *new_list = PyList_New(len);
     if (new_list == NULL) {
         return NULL;
     }
 
+    PyObject *tup0 = PyTuple_New(2);
+    if (tup0 == NULL) {
+        Py_DECREF(new_list);
+        return NULL;
+    }
+
     Py_INCREF(fut->fut_callback0);
-    PyList_SET_ITEM(new_list, 0, fut->fut_callback0);
-    for (i = 0; i < len; i++) {
-        PyObject *cb = PyList_GET_ITEM(fut->fut_callbacks, i);
-        Py_INCREF(cb);
-        PyList_SET_ITEM(new_list, i + 1, cb);
+    PyTuple_SET_ITEM(tup0, 0, fut->fut_callback0);
+    assert(fut->fut_context0 != NULL);
+    Py_INCREF(fut->fut_context0);
+    PyTuple_SET_ITEM(tup0, 1, (PyObject *)fut->fut_context0);
+
+    PyList_SET_ITEM(new_list, 0, tup0);
+
+    if (fut->fut_callbacks != NULL) {
+        for (i = 0; i < PyList_GET_SIZE(fut->fut_callbacks); i++) {
+            PyObject *cb = PyList_GET_ITEM(fut->fut_callbacks, i);
+            Py_INCREF(cb);
+            PyList_SET_ITEM(new_list, i + 1, cb);
+        }
     }
 
     return new_list;
@@ -1912,6 +1981,11 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop)
         return -1;
     }
 
+    self->task_context = PyContext_CopyCurrent();
+    if (self->task_context == NULL) {
+        return -1;
+    }
+
     self->task_fut_waiter = NULL;
     self->task_must_cancel = 0;
     self->task_log_destroy_pending = 1;
@@ -1928,6 +2002,7 @@ static int
 TaskObj_clear(TaskObj *task)
 {
     (void)FutureObj_clear((FutureObj*) task);
+    Py_CLEAR(task->task_context);
     Py_CLEAR(task->task_coro);
     Py_CLEAR(task->task_fut_waiter);
     return 0;
@@ -1936,6 +2011,7 @@ TaskObj_clear(TaskObj *task)
 static int
 TaskObj_traverse(TaskObj *task, visitproc visit, void *arg)
 {
+    Py_VISIT(task->task_context);
     Py_VISIT(task->task_coro);
     Py_VISIT(task->task_fut_waiter);
     (void)FutureObj_traverse((FutureObj*) task, visit, arg);
@@ -2451,7 +2527,7 @@ task_call_step_soon(TaskObj *task, PyObject *arg)
         return -1;
     }
 
-    int ret = call_soon(task->task_loop, cb, NULL);
+    int ret = call_soon(task->task_loop, cb, NULL, task->task_context);
     Py_DECREF(cb);
     return ret;
 }
@@ -2650,7 +2726,8 @@ set_exception:
             if (wrapper == NULL) {
                 goto fail;
             }
-            res = future_add_done_callback((FutureObj*)result, wrapper);
+            res = future_add_done_callback(
+                (FutureObj*)result, wrapper, task->task_context);
             Py_DECREF(wrapper);
             if (res == NULL) {
                 goto fail;
@@ -2724,14 +2801,23 @@ set_exception:
                     goto fail;
                 }
 
-                /* result.add_done_callback(task._wakeup) */
                 wrapper = TaskWakeupMethWrapper_new(task);
                 if (wrapper == NULL) {
                     goto fail;
                 }
-                res = _PyObject_CallMethodIdObjArgs(result,
-                                                    &PyId_add_done_callback,
-                                                    wrapper, NULL);
+
+                /* result.add_done_callback(task._wakeup) */
+                PyObject *add_cb = _PyObject_GetAttrId(
+                    result, &PyId_add_done_callback);
+                if (add_cb == NULL) {
+                    goto fail;
+                }
+                PyObject *stack[2];
+                stack[0] = wrapper;
+                stack[1] = (PyObject *)task->task_context;
+                res = _PyObject_FastCallKeywords(
+                    add_cb, stack, 1, context_kwname);
+                Py_DECREF(add_cb);
                 Py_DECREF(wrapper);
                 if (res == NULL) {
                     goto fail;
@@ -3141,6 +3227,8 @@ module_free(void *m)
     Py_CLEAR(current_tasks);
     Py_CLEAR(iscoroutine_typecache);
 
+    Py_CLEAR(context_kwname);
+
     module_free_freelists();
 }
 
@@ -3164,6 +3252,17 @@ module_init(void)
         goto fail;
     }
 
+
+    context_kwname = PyTuple_New(1);
+    if (context_kwname == NULL) {
+        goto fail;
+    }
+    PyObject *context_str = PyUnicode_FromString("context");
+    if (context_str == NULL) {
+        goto fail;
+    }
+    PyTuple_SET_ITEM(context_kwname, 0, context_str);
+
 #define WITH_MOD(NAME) \
     Py_CLEAR(module); \
     module = PyImport_ImportModule(NAME); \