Add a test for co_extra.

ericsnowcurrently · ericsnowcurrently · commit b943bb0183f3 · 2025-05-06T10:25:55.000-06:00
diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
@@ -1088,6 +1088,18 @@ def test_stateless(self):
                     with self.assertRaises(Exception):
                         _testinternalcapi.verify_stateless_code(func)
 
+        def spam():
+            pass
+
+        with self.subTest('with co_extra'):
+            _testinternalcapi.verify_stateless_code(spam)
+            extra = 'spam'
+            _testinternalcapi.code_set_co_extra(spam.__code__, 0, extra)
+            with self.assertRaises(ValueError):
+                _testinternalcapi.verify_stateless_code(spam)
+            _testinternalcapi.code_set_co_extra(spam.__code__, 0, expect=extra)
+            _testinternalcapi.verify_stateless_code(spam)
+
 
 def isinterned(s):
     return s is sys.intern(('_' + s + '_')[1:-1])
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -1206,6 +1206,58 @@ verify_stateless_code(PyObject *self, PyObject *args, PyObject *kwargs)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+code_set_co_extra(PyObject *self, PyObject *args, PyObject *kwargs)
+{
+    PyObject *codearg;
+    Py_ssize_t index;
+    PyObject *value = NULL;
+    PyObject *expected = NULL;
+    PyObject *notset = Py_None;
+    static char *kwlist[] =
+            {"code", "index", "value", "expect", "notset", NULL};
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs,
+                    "O!n|OOO:code_set_co_extra", kwlist,
+                    &PyCode_Type, &codearg, &index, &value, &expected, &notset))
+    {
+        return NULL;
+    }
+
+    void *extra;
+    if (PyUnstable_Code_GetExtra(codearg, index, &extra) < 0) {
+        return NULL;
+    }
+
+    PyObject *old;
+    if (extra == NULL) {
+        old = Py_NewRef(notset);
+    }
+    else if (extra == expected) {
+        old = Py_NewRef(expected);
+    }
+    else if (extra == value) {
+        old = Py_NewRef(value);
+    }
+    else {
+        if (expected == NULL) {
+            PyErr_SetString(PyExc_ValueError,
+                            "expected existing co_extra to be NULL");
+        }
+        else {
+            PyErr_Format(PyExc_ValueError,
+                         "expected existing co_extra to be %R", expected);
+        }
+        return NULL;
+    }
+
+    if (PyUnstable_Code_SetExtra(codearg, index, value) < 0) {
+        Py_DECREF(old);
+        return NULL;
+    }
+
+    return old;
+}
+
 #ifdef _Py_TIER2
 
 static PyObject *
@@ -2335,6 +2387,8 @@ static PyMethodDef module_functions[] = {
      METH_VARARGS | METH_KEYWORDS, NULL},
     {"verify_stateless_code", _PyCFunction_CAST(verify_stateless_code),
      METH_VARARGS | METH_KEYWORDS, NULL},
+    {"code_set_co_extra", _PyCFunction_CAST(code_set_co_extra),
+     METH_VARARGS | METH_KEYWORDS, NULL},
 #ifdef _Py_TIER2
     {"add_executor_dependency", add_executor_dependency, METH_VARARGS, NULL},
     {"invalidate_executors", invalidate_executors, METH_O, NULL},
