1 files changed, 115 insertions, 0 deletions
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;
+}