Add API function to access code object of generator.

python · markshannon · Feb 23, 2023 · Jan 4, 2023 · Jan 4, 2023 · Jan 4, 2023
commit 3a9df11651c8d2cf5090f26f7e62bd32da6a9813
@@ -44,6 +44,7 @@ PyAPI_FUNC(PyObject *) PyGen_NewWithQualName(PyFrameObject *,
 PyAPI_FUNC(int) _PyGen_SetStopIterationValue(PyObject *);
 PyAPI_FUNC(int) _PyGen_FetchStopIterationValue(PyObject **);
 PyAPI_FUNC(void) _PyGen_Finalize(PyObject *self);
+PyAPI_FUNC(PyCodeObject *) PyGen_GetCode(PyGenObject *gen);
 
 
 /* --- PyCoroObject ------------------------------------------------------- */

@@ -1214,6 +1214,11 @@ def test_pendingcalls_non_threaded(self):
         self.pendingcalls_submit(l, n)
         self.pendingcalls_wait(l, n)
 
+    def test_gen_get_code(self):
+       def genf(): yield
+       gen = genf()
+       self.assertEqual(_testcapi.gen_get_code(gen), gen.gi_code)
+
 
 class SubinterpreterTest(unittest.TestCase):
 

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -2943,6 +2943,16 @@ eval_get_func_desc(PyObject *self, PyObject *func)
     return PyUnicode_FromString(PyEval_GetFuncDesc(func));
 }
 
+static PyObject *
+gen_get_code(PyObject *self, PyObject *gen)
+{
+    if (!PyGen_Check(gen)) {
+        PyErr_SetString(PyExc_TypeError, "argument must be a generator object");
+        return NULL;
+    }
+    return (PyObject *)PyGen_GetCode((PyGenObject *)gen);
+}
+
 static PyObject *
 get_feature_macros(PyObject *self, PyObject *Py_UNUSED(args))
 {
@@ -3346,6 +3356,7 @@ static PyMethodDef TestMethods[] = {
     {"frame_getvarstring", test_frame_getvarstring, METH_VARARGS, NULL},
     {"eval_get_func_name", eval_get_func_name, METH_O, NULL},
     {"eval_get_func_desc", eval_get_func_desc, METH_O, NULL},
+    {"gen_get_code", gen_get_code, METH_O, NULL},
     {"get_feature_macros", get_feature_macros, METH_NOARGS, NULL},
     {"test_code_api", test_code_api, METH_NOARGS, NULL},
     {"settrace_to_record", settrace_to_record, METH_O, NULL},

@@ -31,6 +31,14 @@ _PyGen_GetCode(PyGenObject *gen) {
     return frame->f_code;
 }
 
+PyCodeObject *
+PyGen_GetCode(PyGenObject *gen) {
+    assert(PyGen_Check(gen));
+    PyCodeObject *res = _PyGen_GetCode(gen);
+    Py_INCREF(res);
+    return res;
+}
+
 static inline int
 exc_state_traverse(_PyErr_StackItem *exc_state, visitproc visit, void *arg)
 {