bpo-32455: Add jump parameter to dis.stack_effect(). (GH-6610)

Add C API function PyCompile_OpcodeStackEffectWithJump().

1 files changed, 21 insertions, 3 deletions

diff --git a/Modules/_opcode.c b/Modules/_opcode.c
index f9c1c01..42a8732 100644
--- a/Modules/_opcode.c
+++ b/Modules/_opcode.c
@@ -15,16 +15,20 @@ _opcode.stack_effect -> int
   opcode: int
   oparg: object = None
   /
+  *
+  jump: object = None
 
 Compute the stack effect of the opcode.
 [clinic start generated code]*/
 
 static int
-_opcode_stack_effect_impl(PyObject *module, int opcode, PyObject *oparg)
-/*[clinic end generated code: output=ad39467fa3ad22ce input=2d0a9ee53c0418f5]*/
+_opcode_stack_effect_impl(PyObject *module, int opcode, PyObject *oparg,
+                          PyObject *jump)
+/*[clinic end generated code: output=64a18f2ead954dbb input=461c9d4a44851898]*/
 {
     int effect;
     int oparg_int = 0;
+    int jump_int;
     if (HAS_ARG(opcode)) {
         if (oparg == Py_None) {
             PyErr_SetString(PyExc_ValueError,
@@ -40,7 +44,21 @@ _opcode_stack_effect_impl(PyObject *module, int opcode, PyObject *oparg)
                 "stack_effect: opcode does not permit oparg but oparg was specified");
         return -1;
     }
-    effect = PyCompile_OpcodeStackEffect(opcode, oparg_int);
+    if (jump == Py_None) {
+        jump_int = -1;
+    }
+    else if (jump == Py_True) {
+        jump_int = 1;
+    }
+    else if (jump == Py_False) {
+        jump_int = 0;
+    }
+    else {
+        PyErr_SetString(PyExc_ValueError,
+                "stack_effect: jump must be False, True or None");
+        return -1;
+    }
+    effect = PyCompile_OpcodeStackEffectWithJump(opcode, oparg_int, jump_int);
     if (effect == PY_INVALID_STACK_EFFECT) {
             PyErr_SetString(PyExc_ValueError,
                     "invalid opcode or oparg");