4 files changed, 120 insertions, 1 deletions

diff --git a/Modules/Setup.stdlib.in b/Modules/Setup.stdlib.in
index 7551e5b..b12290d 100644
--- a/Modules/Setup.stdlib.in
+++ b/Modules/Setup.stdlib.in
@@ -169,7 +169,7 @@
 @MODULE__XXTESTFUZZ_TRUE@_xxtestfuzz _xxtestfuzz/_xxtestfuzz.c _xxtestfuzz/fuzzer.c
 @MODULE__TESTBUFFER_TRUE@_testbuffer _testbuffer.c
 @MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c
-@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/unicode.c _testcapi/getargs.c _testcapi/pytime.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/structmember.c _testcapi/exceptions.c
+@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/unicode.c _testcapi/getargs.c _testcapi/pytime.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c
 @MODULE__TESTCLINIC_TRUE@_testclinic _testclinic.c
 
 # Some testing modules MUST be built as shared libraries.
diff --git a/Modules/_testcapi/code.c b/Modules/_testcapi/code.c
new file mode 100644
index 0000000..588dc67
--- /dev/null
+++ b/Modules/_testcapi/code.c
@@ -0,0 +1,115 @@
+#include "parts.h"
+
+static Py_ssize_t
+get_code_extra_index(PyInterpreterState* interp) {
+    Py_ssize_t result = -1;
+
+    static const char *key = "_testcapi.frame_evaluation.code_index";
+
+    PyObject *interp_dict = PyInterpreterState_GetDict(interp); // borrowed
+    assert(interp_dict);  // real users would handle missing dict... somehow
+
+    PyObject *index_obj = PyDict_GetItemString(interp_dict, key); // borrowed
+    Py_ssize_t index = 0;
+    if (!index_obj) {
+        if (PyErr_Occurred()) {
+            goto finally;
+        }
+        index = PyUnstable_Eval_RequestCodeExtraIndex(NULL);
+        if (index < 0 || PyErr_Occurred()) {
+            goto finally;
+        }
+        index_obj = PyLong_FromSsize_t(index); // strong ref
+        if (!index_obj) {
+            goto finally;
+        }
+        int res = PyDict_SetItemString(interp_dict, key, index_obj);
+        Py_DECREF(index_obj);
+        if (res < 0) {
+            goto finally;
+        }
+    }
+    else {
+        index = PyLong_AsSsize_t(index_obj);
+        if (index == -1 && PyErr_Occurred()) {
+            goto finally;
+        }
+    }
+
+    result = index;
+finally:
+    return result;
+}
+
+static PyObject *
+test_code_extra(PyObject* self, PyObject *Py_UNUSED(callable))
+{
+    PyObject *result = NULL;
+    PyObject *test_module = NULL;
+    PyObject *test_func = NULL;
+
+    // Get or initialize interpreter-specific code object storage index
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    if (!interp) {
+        return NULL;
+    }
+    Py_ssize_t code_extra_index = get_code_extra_index(interp);
+    if (PyErr_Occurred()) {
+        goto finally;
+    }
+
+    // Get a function to test with
+    // This can be any Python function. Use `test.test_misc.testfunction`.
+    test_module = PyImport_ImportModule("test.test_capi.test_misc");
+    if (!test_module) {
+        goto finally;
+    }
+    test_func = PyObject_GetAttrString(test_module, "testfunction");
+    if (!test_func) {
+        goto finally;
+    }
+    PyObject *test_func_code = PyFunction_GetCode(test_func);  // borrowed
+    if (!test_func_code) {
+        goto finally;
+    }
+
+    // Check the value is initially NULL
+    void *extra;
+    int res = PyUnstable_Code_GetExtra(test_func_code, code_extra_index, &extra);
+    if (res < 0) {
+        goto finally;
+    }
+    assert (extra == NULL);
+
+    // Set another code extra value
+    res = PyUnstable_Code_SetExtra(test_func_code, code_extra_index, (void*)(uintptr_t)77);
+    if (res < 0) {
+        goto finally;
+    }
+    // Assert it was set correctly
+    res = PyUnstable_Code_GetExtra(test_func_code, code_extra_index, &extra);
+    if (res < 0) {
+        goto finally;
+    }
+    assert ((uintptr_t)extra == 77);
+
+    result = Py_NewRef(Py_None);
+finally:
+    Py_XDECREF(test_module);
+    Py_XDECREF(test_func);
+    return result;
+}
+
+static PyMethodDef TestMethods[] = {
+    {"test_code_extra", test_code_extra, METH_NOARGS},
+    {NULL},
+};
+
+int
+_PyTestCapi_Init_Code(PyObject *m) {
+    if (PyModule_AddFunctions(m, TestMethods) < 0) {
+        return -1;
+    }
+
+    return 0;
+}
diff --git a/Modules/_testcapi/parts.h b/Modules/_testcapi/parts.h
index 1689f18..c8f31dc 100644
--- a/Modules/_testcapi/parts.h
+++ b/Modules/_testcapi/parts.h
@@ -37,6 +37,7 @@ int _PyTestCapi_Init_Long(PyObject *module);
 int _PyTestCapi_Init_Float(PyObject *module);
 int _PyTestCapi_Init_Structmember(PyObject *module);
 int _PyTestCapi_Init_Exceptions(PyObject *module);
+int _PyTestCapi_Init_Code(PyObject *module);
 
 #ifdef LIMITED_API_AVAILABLE
 int _PyTestCapi_Init_VectorcallLimited(PyObject *module);
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index fc716a3..10e507d 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -4083,6 +4083,9 @@ PyInit__testcapi(void)
     if (_PyTestCapi_Init_Exceptions(m) < 0) {
         return NULL;
     }
+    if (_PyTestCapi_Init_Code(m) < 0) {
+        return NULL;
+    }
 
 #ifndef LIMITED_API_AVAILABLE
     PyModule_AddObjectRef(m, "LIMITED_API_AVAILABLE", Py_False);