1 #include "parts.h"
2 #include "util.h"
3
4 static Py_ssize_t
get_code_extra_index(PyInterpreterState * interp)5 get_code_extra_index(PyInterpreterState* interp) {
6 Py_ssize_t result = -1;
7
8 static const char *key = "_testcapi.frame_evaluation.code_index";
9
10 PyObject *interp_dict = PyInterpreterState_GetDict(interp); // borrowed
11 assert(interp_dict); // real users would handle missing dict... somehow
12
13 PyObject *index_obj;
14 if (PyDict_GetItemStringRef(interp_dict, key, &index_obj) < 0) {
15 goto finally;
16 }
17 Py_ssize_t index = 0;
18 if (!index_obj) {
19 index = PyUnstable_Eval_RequestCodeExtraIndex(NULL);
20 if (index < 0 || PyErr_Occurred()) {
21 goto finally;
22 }
23 index_obj = PyLong_FromSsize_t(index); // strong ref
24 if (!index_obj) {
25 goto finally;
26 }
27 int res = PyDict_SetItemString(interp_dict, key, index_obj);
28 Py_DECREF(index_obj);
29 if (res < 0) {
30 goto finally;
31 }
32 }
33 else {
34 index = PyLong_AsSsize_t(index_obj);
35 Py_DECREF(index_obj);
36 if (index == -1 && PyErr_Occurred()) {
37 goto finally;
38 }
39 }
40
41 result = index;
42 finally:
43 return result;
44 }
45
46 static PyObject *
test_code_extra(PyObject * self,PyObject * Py_UNUSED (callable))47 test_code_extra(PyObject* self, PyObject *Py_UNUSED(callable))
48 {
49 PyObject *result = NULL;
50 PyObject *test_module = NULL;
51 PyObject *test_func = NULL;
52
53 // Get or initialize interpreter-specific code object storage index
54 PyInterpreterState *interp = PyInterpreterState_Get();
55 if (!interp) {
56 return NULL;
57 }
58 Py_ssize_t code_extra_index = get_code_extra_index(interp);
59 if (PyErr_Occurred()) {
60 goto finally;
61 }
62
63 // Get a function to test with
64 // This can be any Python function. Use `test.test_misc.testfunction`.
65 test_module = PyImport_ImportModule("test.test_capi.test_misc");
66 if (!test_module) {
67 goto finally;
68 }
69 test_func = PyObject_GetAttrString(test_module, "testfunction");
70 if (!test_func) {
71 goto finally;
72 }
73 PyObject *test_func_code = PyFunction_GetCode(test_func); // borrowed
74 if (!test_func_code) {
75 goto finally;
76 }
77
78 // Check the value is initially NULL
79 void *extra = UNINITIALIZED_PTR;
80 int res = PyUnstable_Code_GetExtra(test_func_code, code_extra_index, &extra);
81 if (res < 0) {
82 goto finally;
83 }
84 assert (extra == NULL);
85
86 // Set another code extra value
87 res = PyUnstable_Code_SetExtra(test_func_code, code_extra_index, (void*)(uintptr_t)77);
88 if (res < 0) {
89 goto finally;
90 }
91 // Assert it was set correctly
92 extra = UNINITIALIZED_PTR;
93 res = PyUnstable_Code_GetExtra(test_func_code, code_extra_index, &extra);
94 if (res < 0) {
95 goto finally;
96 }
97 assert ((uintptr_t)extra == 77);
98 // Revert to initial code extra value.
99 res = PyUnstable_Code_SetExtra(test_func_code, code_extra_index, NULL);
100 if (res < 0) {
101 goto finally;
102 }
103 result = Py_NewRef(Py_None);
104 finally:
105 Py_XDECREF(test_module);
106 Py_XDECREF(test_func);
107 return result;
108 }
109
110 static PyMethodDef TestMethods[] = {
111 {"test_code_extra", test_code_extra, METH_NOARGS},
112 {NULL},
113 };
114
115 int
_PyTestCapi_Init_Code(PyObject * m)116 _PyTestCapi_Init_Code(PyObject *m) {
117 if (PyModule_AddFunctions(m, TestMethods) < 0) {
118 return -1;
119 }
120
121 return 0;
122 }
123