1  #include "parts.h"
       2  #include "util.h"
       3  
       4  static Py_ssize_t
       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 = _PyDict_GetItemStringWithError(interp_dict, key); // borrowed
      14      Py_ssize_t index = 0;
      15      if (!index_obj) {
      16          if (PyErr_Occurred()) {
      17              goto finally;
      18          }
      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          if (index == -1 && PyErr_Occurred()) {
      36              goto finally;
      37          }
      38      }
      39  
      40      result = index;
      41  finally:
      42      return result;
      43  }
      44  
      45  static PyObject *
      46  test_code_extra(PyObject* self, PyObject *Py_UNUSED(callable))
      47  {
      48      PyObject *result = NULL;
      49      PyObject *test_module = NULL;
      50      PyObject *test_func = NULL;
      51  
      52      // Get or initialize interpreter-specific code object storage index
      53      PyInterpreterState *interp = PyInterpreterState_Get();
      54      if (!interp) {
      55          return NULL;
      56      }
      57      Py_ssize_t code_extra_index = get_code_extra_index(interp);
      58      if (PyErr_Occurred()) {
      59          goto finally;
      60      }
      61  
      62      // Get a function to test with
      63      // This can be any Python function. Use `test.test_misc.testfunction`.
      64      test_module = PyImport_ImportModule("test.test_capi.test_misc");
      65      if (!test_module) {
      66          goto finally;
      67      }
      68      test_func = PyObject_GetAttrString(test_module, "testfunction");
      69      if (!test_func) {
      70          goto finally;
      71      }
      72      PyObject *test_func_code = PyFunction_GetCode(test_func);  // borrowed
      73      if (!test_func_code) {
      74          goto finally;
      75      }
      76  
      77      // Check the value is initially NULL
      78      void *extra = UNINITIALIZED_PTR;
      79      int res = PyUnstable_Code_GetExtra(test_func_code, code_extra_index, &extra);
      80      if (res < 0) {
      81          goto finally;
      82      }
      83      assert (extra == NULL);
      84  
      85      // Set another code extra value
      86      res = PyUnstable_Code_SetExtra(test_func_code, code_extra_index, (void*)(uintptr_t)77);
      87      if (res < 0) {
      88          goto finally;
      89      }
      90      // Assert it was set correctly
      91      extra = UNINITIALIZED_PTR;
      92      res = PyUnstable_Code_GetExtra(test_func_code, code_extra_index, &extra);
      93      if (res < 0) {
      94          goto finally;
      95      }
      96      assert ((uintptr_t)extra == 77);
      97      // Revert to initial code extra value.
      98      res = PyUnstable_Code_SetExtra(test_func_code, code_extra_index, NULL);
      99      if (res < 0) {
     100          goto finally;
     101      }
     102      result = Py_NewRef(Py_None);
     103  finally:
     104      Py_XDECREF(test_module);
     105      Py_XDECREF(test_func);
     106      return result;
     107  }
     108  
     109  static PyMethodDef TestMethods[] = {
     110      {"test_code_extra", test_code_extra, METH_NOARGS},
     111      {NULL},
     112  };
     113  
     114  int
     115  _PyTestCapi_Init_Code(PyObject *m) {
     116      if (PyModule_AddFunctions(m, TestMethods) < 0) {
     117          return -1;
     118      }
     119  
     120      return 0;
     121  }