bpo-32436: Implement PEP 567 (#5027)

7 files changed, 287 insertions, 58 deletions

diff --git a/Modules/Setup.dist b/Modules/Setup.dist
index 1f2d56c..239550e 100644
--- a/Modules/Setup.dist
+++ b/Modules/Setup.dist
@@ -176,6 +176,7 @@ _symtable symtablemodule.c
 #array arraymodule.c	# array objects
 #cmath cmathmodule.c _math.c # -lm # complex math library functions
 #math mathmodule.c _math.c # -lm # math library functions, e.g. sin()
+#_contextvars _contextvarsmodule.c  # Context Variables
 #_struct _struct.c	# binary structure packing/unpacking
 #_weakref _weakref.c	# basic weak reference support
 #_testcapi _testcapimodule.c    # Python C API test module
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); \
diff --git a/Modules/_contextvarsmodule.c b/Modules/_contextvarsmodule.c
new file mode 100644
index 0000000..b7d112d
--- /dev/null
+++ b/Modules/_contextvarsmodule.c
@@ -0,0 +1,75 @@
+#include "Python.h"
+
+#include "clinic/_contextvarsmodule.c.h"
+
+/*[clinic input]
+module _contextvars
+[clinic start generated code]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=a0955718c8b8cea6]*/
+
+
+/*[clinic input]
+_contextvars.copy_context
+[clinic start generated code]*/
+
+static PyObject *
+_contextvars_copy_context_impl(PyObject *module)
+/*[clinic end generated code: output=1fcd5da7225c4fa9 input=89bb9ae485888440]*/
+{
+    return (PyObject *)PyContext_CopyCurrent();
+}
+
+
+PyDoc_STRVAR(module_doc, "Context Variables");
+
+static PyMethodDef _contextvars_methods[] = {
+    _CONTEXTVARS_COPY_CONTEXT_METHODDEF
+    {NULL, NULL}
+};
+
+static struct PyModuleDef _contextvarsmodule = {
+    PyModuleDef_HEAD_INIT,      /* m_base */
+    "_contextvars",             /* m_name */
+    module_doc,                 /* m_doc */
+    -1,                         /* m_size */
+    _contextvars_methods,       /* m_methods */
+    NULL,                       /* m_slots */
+    NULL,                       /* m_traverse */
+    NULL,                       /* m_clear */
+    NULL,                       /* m_free */
+};
+
+PyMODINIT_FUNC
+PyInit__contextvars(void)
+{
+    PyObject *m = PyModule_Create(&_contextvarsmodule);
+    if (m == NULL) {
+        return NULL;
+    }
+
+    Py_INCREF(&PyContext_Type);
+    if (PyModule_AddObject(m, "Context",
+                           (PyObject *)&PyContext_Type) < 0)
+    {
+        Py_DECREF(&PyContext_Type);
+        return NULL;
+    }
+
+    Py_INCREF(&PyContextVar_Type);
+    if (PyModule_AddObject(m, "ContextVar",
+                           (PyObject *)&PyContextVar_Type) < 0)
+    {
+        Py_DECREF(&PyContextVar_Type);
+        return NULL;
+    }
+
+    Py_INCREF(&PyContextToken_Type);
+    if (PyModule_AddObject(m, "Token",
+                           (PyObject *)&PyContextToken_Type) < 0)
+    {
+        Py_DECREF(&PyContextToken_Type);
+        return NULL;
+    }
+
+    return m;
+}
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 0d6bf45..e3be7d3 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -4438,6 +4438,13 @@ test_pythread_tss_key_state(PyObject *self, PyObject *args)
 }
 
 
+static PyObject*
+new_hamt(PyObject *self, PyObject *args)
+{
+    return _PyContext_NewHamtForTests();
+}
+
+
 static PyMethodDef TestMethods[] = {
     {"raise_exception",         raise_exception,                 METH_VARARGS},
     {"raise_memoryerror",   (PyCFunction)raise_memoryerror,  METH_NOARGS},
@@ -4655,6 +4662,7 @@ static PyMethodDef TestMethods[] = {
     {"get_mapping_values", get_mapping_values, METH_O},
     {"get_mapping_items", get_mapping_items, METH_O},
     {"test_pythread_tss_key_state", test_pythread_tss_key_state, METH_VARARGS},
+    {"hamt", new_hamt, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h
index f2e0f40..9fc9d6b 100644
--- a/Modules/clinic/_asynciomodule.c.h
+++ b/Modules/clinic/_asynciomodule.c.h
@@ -110,7 +110,7 @@ PyDoc_STRVAR(_asyncio_Future_set_exception__doc__,
     {"set_exception", (PyCFunction)_asyncio_Future_set_exception, METH_O, _asyncio_Future_set_exception__doc__},
 
 PyDoc_STRVAR(_asyncio_Future_add_done_callback__doc__,
-"add_done_callback($self, fn, /)\n"
+"add_done_callback($self, fn, /, *, context=None)\n"
 "--\n"
 "\n"
 "Add a callback to be run when the future becomes done.\n"
@@ -120,7 +120,30 @@ PyDoc_STRVAR(_asyncio_Future_add_done_callback__doc__,
 "scheduled with call_soon.");
 
 #define _ASYNCIO_FUTURE_ADD_DONE_CALLBACK_METHODDEF    \
-    {"add_done_callback", (PyCFunction)_asyncio_Future_add_done_callback, METH_O, _asyncio_Future_add_done_callback__doc__},
+    {"add_done_callback", (PyCFunction)_asyncio_Future_add_done_callback, METH_FASTCALL|METH_KEYWORDS, _asyncio_Future_add_done_callback__doc__},
+
+static PyObject *
+_asyncio_Future_add_done_callback_impl(FutureObj *self, PyObject *fn,
+                                       PyObject *context);
+
+static PyObject *
+_asyncio_Future_add_done_callback(FutureObj *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    static const char * const _keywords[] = {"", "context", NULL};
+    static _PyArg_Parser _parser = {"O|$O:add_done_callback", _keywords, 0};
+    PyObject *fn;
+    PyObject *context = NULL;
+
+    if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
+        &fn, &context)) {
+        goto exit;
+    }
+    return_value = _asyncio_Future_add_done_callback_impl(self, fn, context);
+
+exit:
+    return return_value;
+}
 
 PyDoc_STRVAR(_asyncio_Future_remove_done_callback__doc__,
 "remove_done_callback($self, fn, /)\n"
@@ -763,4 +786,4 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=616e814431893dcc input=a9049054013a1b77]*/
+/*[clinic end generated code: output=bcbaf1b2480f4aa9 input=a9049054013a1b77]*/
diff --git a/Modules/clinic/_contextvarsmodule.c.h b/Modules/clinic/_contextvarsmodule.c.h
new file mode 100644
index 0000000..b1885e4
--- /dev/null
+++ b/Modules/clinic/_contextvarsmodule.c.h
@@ -0,0 +1,21 @@
+/*[clinic input]
+preserve
+[clinic start generated code]*/
+
+PyDoc_STRVAR(_contextvars_copy_context__doc__,
+"copy_context($module, /)\n"
+"--\n"
+"\n");
+
+#define _CONTEXTVARS_COPY_CONTEXT_METHODDEF    \
+    {"copy_context", (PyCFunction)_contextvars_copy_context, METH_NOARGS, _contextvars_copy_context__doc__},
+
+static PyObject *
+_contextvars_copy_context_impl(PyObject *module);
+
+static PyObject *
+_contextvars_copy_context(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    return _contextvars_copy_context_impl(module);
+}
+/*[clinic end generated code: output=26e07024451baf52 input=a9049054013a1b77]*/
diff --git a/Modules/gcmodule.c b/Modules/gcmodule.c
index ea3c294..8ba1093 100644
--- a/Modules/gcmodule.c
+++ b/Modules/gcmodule.c
@@ -24,6 +24,7 @@
 */
 
 #include "Python.h"
+#include "internal/context.h"
 #include "internal/mem.h"
 #include "internal/pystate.h"
 #include "frameobject.h"        /* for PyFrame_ClearFreeList */
@@ -790,6 +791,7 @@ clear_freelists(void)
     (void)PyDict_ClearFreeList();
     (void)PySet_ClearFreeList();
     (void)PyAsyncGen_ClearFreeLists();
+    (void)PyContext_ClearFreeList();
 }
 
 /* This is the main function.  Read this to understand how the