Add API function to access code object of generator.

markshannon · markshannon · commit 3a9df11651c8 · 2023-01-04T15:31:32.000Z
diff --git a/Include/cpython/genobject.h b/Include/cpython/genobject.h
@@ -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 ------------------------------------------------------- */
diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py
@@ -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},
diff --git a/Objects/genobject.c b/Objects/genobject.c
@@ -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)
 {

Original file line number	Diff line number	Diff line change
`@@ -31,6 +31,14 @@ _PyGen_GetCode(PyGenObject *gen) {`
`31`	`31`	`return frame->f_code;`
`32`	`32`	`}`
`33`	`33`
	`34`	`+PyCodeObject *`
	`35`	`+PyGen_GetCode(PyGenObject *gen) {`
	`36`	`+ assert(PyGen_Check(gen));`
	`37`	`+ PyCodeObject *res = _PyGen_GetCode(gen);`
	`38`	`+ Py_INCREF(res);`
	`39`	`+ return res;`
	`40`	`+}`
	`41`	`+`
`34`	`42`	`static inline int`
`35`	`43`	`exc_state_traverse(_PyErr_StackItem exc_state, visitproc visit, void arg)`
`36`	`44`	`{`