From da3f85480437683175461f7124c5ae14b7794586 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Wed, 25 Apr 2018 23:17:44 +0300
Subject: [PATCH 1/4] bpo-32455: Add jump parameter to dis.stack_effect().

Add jump parameter to undocumented C API function
PyCompile_OpcodeStackEffect().
---
 Doc/library/dis.rst                           | 11 +++-
 Include/compile.h                             |  2 +-
 Lib/test/test__opcode.py                      | 61 ++++++++++++++-----
 .../2018-04-26-13-31-10.bpo-32455.KPWg3K.rst  |  2 +
 Modules/_opcode.c                             | 23 ++++++-
 Modules/clinic/_opcode.c.h                    | 20 +++---
 Python/compile.c                              |  4 +-
 7 files changed, 94 insertions(+), 29 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst

diff --git a/Doc/library/dis.rst b/Doc/library/dis.rst
index 8f505e65bd51e5b..ef5775085d8b427 100644
--- a/Doc/library/dis.rst
+++ b/Doc/library/dis.rst
@@ -248,12 +248,21 @@ operation is being performed, so the intermediate analysis object isn't useful:
    return a list of these offsets.
 
 
-.. function:: stack_effect(opcode, [oparg])
+.. function:: stack_effect(opcode, oparg=None, *, jump=None)
 
    Compute the stack effect of *opcode* with argument *oparg*.
 
+   If the code has a jump target and *jump* is ``True``, :func:`~stack_effect`
+   will return the stack effect of jumping.  If *jump* is ``False``,
+   it will return the stack effect of not jumping. And if *jump* is
+   ``None`` (default), it will return the maximal stack effect of both cases.
+
    .. versionadded:: 3.4
 
+   .. versionchanged:: 3.8
+      Added *jump* parameter.
+
+
 .. _bytecodes:
 
 Python Bytecode Instructions
diff --git a/Include/compile.h b/Include/compile.h
index edb961f4d72042f..629ad25c29ab2a5 100644
--- a/Include/compile.h
+++ b/Include/compile.h
@@ -74,7 +74,7 @@ PyAPI_FUNC(PyFutureFeatures *) PyFuture_FromASTObject(
 PyAPI_FUNC(PyObject*) _Py_Mangle(PyObject *p, PyObject *name);
 
 #define PY_INVALID_STACK_EFFECT INT_MAX
-PyAPI_FUNC(int) PyCompile_OpcodeStackEffect(int opcode, int oparg);
+PyAPI_FUNC(int) PyCompile_OpcodeStackEffect(int opcode, int oparg, int jump);
 
 PyAPI_FUNC(int) _PyAST_Optimize(struct _mod *, PyArena *arena, int optimize);
 
diff --git a/Lib/test/test__opcode.py b/Lib/test/test__opcode.py
index 2af1ee35bff0abb..0fb39eed6064251 100644
--- a/Lib/test/test__opcode.py
+++ b/Lib/test/test__opcode.py
@@ -3,32 +3,65 @@
 import unittest
 
 _opcode = import_module("_opcode")
+from _opcode import stack_effect
+
 
 class OpcodeTests(unittest.TestCase):
 
     def test_stack_effect(self):
-        self.assertEqual(_opcode.stack_effect(dis.opmap['POP_TOP']), -1)
-        self.assertEqual(_opcode.stack_effect(dis.opmap['DUP_TOP_TWO']), 2)
-        self.assertEqual(_opcode.stack_effect(dis.opmap['BUILD_SLICE'], 0), -1)
-        self.assertEqual(_opcode.stack_effect(dis.opmap['BUILD_SLICE'], 1), -1)
-        self.assertEqual(_opcode.stack_effect(dis.opmap['BUILD_SLICE'], 3), -2)
-        self.assertRaises(ValueError, _opcode.stack_effect, 30000)
-        self.assertRaises(ValueError, _opcode.stack_effect, dis.opmap['BUILD_SLICE'])
-        self.assertRaises(ValueError, _opcode.stack_effect, dis.opmap['POP_TOP'], 0)
+        self.assertEqual(stack_effect(dis.opmap['POP_TOP']), -1)
+        self.assertEqual(stack_effect(dis.opmap['DUP_TOP_TWO']), 2)
+        self.assertEqual(stack_effect(dis.opmap['BUILD_SLICE'], 0), -1)
+        self.assertEqual(stack_effect(dis.opmap['BUILD_SLICE'], 1), -1)
+        self.assertEqual(stack_effect(dis.opmap['BUILD_SLICE'], 3), -2)
+        self.assertRaises(ValueError, stack_effect, 30000)
+        self.assertRaises(ValueError, stack_effect, dis.opmap['BUILD_SLICE'])
+        self.assertRaises(ValueError, stack_effect, dis.opmap['POP_TOP'], 0)
         # All defined opcodes
         for name, code in dis.opmap.items():
             with self.subTest(opname=name):
                 if code < dis.HAVE_ARGUMENT:
-                    _opcode.stack_effect(code)
-                    self.assertRaises(ValueError, _opcode.stack_effect, code, 0)
+                    stack_effect(code)
+                    self.assertRaises(ValueError, stack_effect, code, 0)
                 else:
-                    _opcode.stack_effect(code, 0)
-                    self.assertRaises(ValueError, _opcode.stack_effect, code)
+                    stack_effect(code, 0)
+                    self.assertRaises(ValueError, stack_effect, code)
         # All not defined opcodes
         for code in set(range(256)) - set(dis.opmap.values()):
             with self.subTest(opcode=code):
-                self.assertRaises(ValueError, _opcode.stack_effect, code)
-                self.assertRaises(ValueError, _opcode.stack_effect, code, 0)
+                self.assertRaises(ValueError, stack_effect, code)
+                self.assertRaises(ValueError, stack_effect, code, 0)
+
+    def test_stack_effect_jump(self):
+        JUMP_IF_TRUE_OR_POP = dis.opmap['JUMP_IF_TRUE_OR_POP']
+        self.assertEqual(stack_effect(JUMP_IF_TRUE_OR_POP, 0), 0)
+        self.assertEqual(stack_effect(JUMP_IF_TRUE_OR_POP, 0, jump=True), 0)
+        self.assertEqual(stack_effect(JUMP_IF_TRUE_OR_POP, 0, jump=False), -1)
+        FOR_ITER = dis.opmap['FOR_ITER']
+        self.assertEqual(stack_effect(FOR_ITER, 0), 1)
+        self.assertEqual(stack_effect(FOR_ITER, 0, jump=True), -1)
+        self.assertEqual(stack_effect(FOR_ITER, 0, jump=False), 1)
+        JUMP_FORWARD = dis.opmap['JUMP_FORWARD']
+        self.assertEqual(stack_effect(JUMP_FORWARD, 0), 0)
+        self.assertEqual(stack_effect(JUMP_FORWARD, 0, jump=True), 0)
+        self.assertEqual(stack_effect(JUMP_FORWARD, 0, jump=False), 0)
+        # All defined opcodes
+        has_jump = dis.hasjabs + dis.hasjrel
+        for name, code in dis.opmap.items():
+            with self.subTest(opname=name):
+                if code < dis.HAVE_ARGUMENT:
+                    common = stack_effect(code)
+                    jump = stack_effect(code, jump=True)
+                    nojump = stack_effect(code, jump=False)
+                else:
+                    common = stack_effect(code, 0)
+                    jump = stack_effect(code, 0, jump=True)
+                    nojump = stack_effect(code, 0, jump=False)
+                if code in has_jump:
+                    self.assertEqual(common, max(jump, nojump))
+                else:
+                    self.assertEqual(jump, common)
+                    self.assertEqual(nojump, common)
 
 
 if __name__ == "__main__":
diff --git a/Misc/NEWS.d/next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst b/Misc/NEWS.d/next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst
new file mode 100644
index 000000000000000..b8ff93818673959
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst
@@ -0,0 +1,2 @@
+Add *jump* parameter to :func:`dis.stack_effect` and to undocumented C API
+function :c:func:`PyCompile_OpcodeStackEffect`.
diff --git a/Modules/_opcode.c b/Modules/_opcode.c
index f9c1c0108d6af72..c278d46f7bac451 100644
--- a/Modules/_opcode.c
+++ b/Modules/_opcode.c
@@ -15,16 +15,19 @@ _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=d5e3992a62215be5]*/
 {
     int effect;
     int oparg_int = 0;
+    int jump_int;
     if (HAS_ARG(opcode)) {
         if (oparg == Py_None) {
             PyErr_SetString(PyExc_ValueError,
@@ -40,7 +43,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_OpcodeStackEffect(opcode, oparg_int, jump_int);
     if (effect == PY_INVALID_STACK_EFFECT) {
             PyErr_SetString(PyExc_ValueError,
                     "invalid opcode or oparg");
diff --git a/Modules/clinic/_opcode.c.h b/Modules/clinic/_opcode.c.h
index 4d593edfac088c4..79cb2a326d4ffc6 100644
--- a/Modules/clinic/_opcode.c.h
+++ b/Modules/clinic/_opcode.c.h
@@ -3,30 +3,34 @@ preserve
 [clinic start generated code]*/
 
 PyDoc_STRVAR(_opcode_stack_effect__doc__,
-"stack_effect($module, opcode, oparg=None, /)\n"
+"stack_effect($module, opcode, oparg=None, /, jump=None)\n"
 "--\n"
 "\n"
 "Compute the stack effect of the opcode.");
 
 #define _OPCODE_STACK_EFFECT_METHODDEF    \
-    {"stack_effect", (PyCFunction)_opcode_stack_effect, METH_FASTCALL, _opcode_stack_effect__doc__},
+    {"stack_effect", (PyCFunction)_opcode_stack_effect, METH_FASTCALL|METH_KEYWORDS, _opcode_stack_effect__doc__},
 
 static int
-_opcode_stack_effect_impl(PyObject *module, int opcode, PyObject *oparg);
+_opcode_stack_effect_impl(PyObject *module, int opcode, PyObject *oparg,
+                          PyObject *jump);
 
 static PyObject *
-_opcode_stack_effect(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
+_opcode_stack_effect(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
 {
     PyObject *return_value = NULL;
+    static const char * const _keywords[] = {"", "", "jump", NULL};
+    static _PyArg_Parser _parser = {"i|OO:stack_effect", _keywords, 0};
     int opcode;
     PyObject *oparg = Py_None;
+    PyObject *jump = Py_None;
     int _return_value;
 
-    if (!_PyArg_ParseStack(args, nargs, "i|O:stack_effect",
-        &opcode, &oparg)) {
+    if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
+        &opcode, &oparg, &jump)) {
         goto exit;
     }
-    _return_value = _opcode_stack_effect_impl(module, opcode, oparg);
+    _return_value = _opcode_stack_effect_impl(module, opcode, oparg, jump);
     if ((_return_value == -1) && PyErr_Occurred()) {
         goto exit;
     }
@@ -35,4 +39,4 @@ _opcode_stack_effect(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=577a91c9aa5559a9 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=4773d6a4221eb554 input=a9049054013a1b77]*/
diff --git a/Python/compile.c b/Python/compile.c
index 3528670ef67503c..5dd354a50c378a9 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1117,9 +1117,9 @@ stack_effect(int opcode, int oparg, int jump)
 }
 
 int
-PyCompile_OpcodeStackEffect(int opcode, int oparg)
+PyCompile_OpcodeStackEffect(int opcode, int oparg, int jump)
 {
-    return stack_effect(opcode, oparg, -1);
+    return stack_effect(opcode, oparg, jump);
 }
 
 /* Add an opcode with no argument.

From 941647a1d70e8d25f1a818bee4f75354f294b032 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Sun, 8 Jul 2018 12:01:51 +0300
Subject: [PATCH 2/4] Keep an old PyCompile_OpcodeStackEffect() and add
 PyCompile_OpcodeStackEffectWithJump().

---
 Include/compile.h | 3 ++-
 Modules/_opcode.c | 2 +-
 Python/compile.c  | 8 +++++++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/Include/compile.h b/Include/compile.h
index 629ad25c29ab2a5..2dacfff37f8c2a5 100644
--- a/Include/compile.h
+++ b/Include/compile.h
@@ -74,7 +74,8 @@ PyAPI_FUNC(PyFutureFeatures *) PyFuture_FromASTObject(
 PyAPI_FUNC(PyObject*) _Py_Mangle(PyObject *p, PyObject *name);
 
 #define PY_INVALID_STACK_EFFECT INT_MAX
-PyAPI_FUNC(int) PyCompile_OpcodeStackEffect(int opcode, int oparg, int jump);
+PyAPI_FUNC(int) PyCompile_OpcodeStackEffect(int opcode, int oparg);
+PyAPI_FUNC(int) PyCompile_OpcodeStackEffectWithJump(int opcode, int oparg, int jump);
 
 PyAPI_FUNC(int) _PyAST_Optimize(struct _mod *, PyArena *arena, int optimize);
 
diff --git a/Modules/_opcode.c b/Modules/_opcode.c
index c278d46f7bac451..116f3d37ed41525 100644
--- a/Modules/_opcode.c
+++ b/Modules/_opcode.c
@@ -57,7 +57,7 @@ _opcode_stack_effect_impl(PyObject *module, int opcode, PyObject *oparg,
                 "stack_effect: jump must be False, True or None");
         return -1;
     }
-    effect = PyCompile_OpcodeStackEffect(opcode, oparg_int, jump_int);
+    effect = PyCompile_OpcodeStackEffectWithJump(opcode, oparg_int, jump_int);
     if (effect == PY_INVALID_STACK_EFFECT) {
             PyErr_SetString(PyExc_ValueError,
                     "invalid opcode or oparg");
diff --git a/Python/compile.c b/Python/compile.c
index 5dd354a50c378a9..da5db608b0b3f9e 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1117,11 +1117,17 @@ stack_effect(int opcode, int oparg, int jump)
 }
 
 int
-PyCompile_OpcodeStackEffect(int opcode, int oparg, int jump)
+PyCompile_OpcodeStackEffectWithJump(int opcode, int oparg, int jump)
 {
     return stack_effect(opcode, oparg, jump);
 }
 
+int
+PyCompile_OpcodeStackEffect(int opcode, int oparg)
+{
+    return stack_effect(opcode, oparg, -1);
+}
+
 /* Add an opcode with no argument.
    Returns 0 on failure, 1 on success.
 */

From 564e96757c5e8ba35231f43e3abf3f6991452156 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Sun, 8 Jul 2018 12:03:13 +0300
Subject: [PATCH 3/4] Make jump a keyword-only parameter.

---
 Modules/_opcode.c          | 3 ++-
 Modules/clinic/_opcode.c.h | 6 +++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/Modules/_opcode.c b/Modules/_opcode.c
index 116f3d37ed41525..42a8732694afef6 100644
--- a/Modules/_opcode.c
+++ b/Modules/_opcode.c
@@ -15,6 +15,7 @@ _opcode.stack_effect -> int
   opcode: int
   oparg: object = None
   /
+  *
   jump: object = None
 
 Compute the stack effect of the opcode.
@@ -23,7 +24,7 @@ Compute the stack effect of the opcode.
 static int
 _opcode_stack_effect_impl(PyObject *module, int opcode, PyObject *oparg,
                           PyObject *jump)
-/*[clinic end generated code: output=64a18f2ead954dbb input=d5e3992a62215be5]*/
+/*[clinic end generated code: output=64a18f2ead954dbb input=461c9d4a44851898]*/
 {
     int effect;
     int oparg_int = 0;
diff --git a/Modules/clinic/_opcode.c.h b/Modules/clinic/_opcode.c.h
index 79cb2a326d4ffc6..b162d84e5db41b9 100644
--- a/Modules/clinic/_opcode.c.h
+++ b/Modules/clinic/_opcode.c.h
@@ -3,7 +3,7 @@ preserve
 [clinic start generated code]*/
 
 PyDoc_STRVAR(_opcode_stack_effect__doc__,
-"stack_effect($module, opcode, oparg=None, /, jump=None)\n"
+"stack_effect($module, opcode, oparg=None, /, *, jump=None)\n"
 "--\n"
 "\n"
 "Compute the stack effect of the opcode.");
@@ -20,7 +20,7 @@ _opcode_stack_effect(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
 {
     PyObject *return_value = NULL;
     static const char * const _keywords[] = {"", "", "jump", NULL};
-    static _PyArg_Parser _parser = {"i|OO:stack_effect", _keywords, 0};
+    static _PyArg_Parser _parser = {"i|O$O:stack_effect", _keywords, 0};
     int opcode;
     PyObject *oparg = Py_None;
     PyObject *jump = Py_None;
@@ -39,4 +39,4 @@ _opcode_stack_effect(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=4773d6a4221eb554 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=bbf6c4cfc91edc29 input=a9049054013a1b77]*/

From 66376aa1eb65da1e7282c0ed1704be043ab9f840 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Sun, 8 Jul 2018 12:06:26 +0300
Subject: [PATCH 4/4] Update news entries.

---
 .../NEWS.d/next/C API/2018-07-08-12-06-18.bpo-32455.KVHlkz.rst | 1 +
 .../next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst      | 3 +--
 2 files changed, 2 insertions(+), 2 deletions(-)
 create mode 100644 Misc/NEWS.d/next/C API/2018-07-08-12-06-18.bpo-32455.KVHlkz.rst

diff --git a/Misc/NEWS.d/next/C API/2018-07-08-12-06-18.bpo-32455.KVHlkz.rst b/Misc/NEWS.d/next/C API/2018-07-08-12-06-18.bpo-32455.KVHlkz.rst
new file mode 100644
index 000000000000000..f28be876cf3917d
--- /dev/null
+++ b/Misc/NEWS.d/next/C API/2018-07-08-12-06-18.bpo-32455.KVHlkz.rst	
@@ -0,0 +1 @@
+Added :c:func:`PyCompile_OpcodeStackEffectWithJump`.
diff --git a/Misc/NEWS.d/next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst b/Misc/NEWS.d/next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst
index b8ff93818673959..dd873f77bbd412e 100644
--- a/Misc/NEWS.d/next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst
+++ b/Misc/NEWS.d/next/Library/2018-04-26-13-31-10.bpo-32455.KPWg3K.rst
@@ -1,2 +1 @@
-Add *jump* parameter to :func:`dis.stack_effect` and to undocumented C API
-function :c:func:`PyCompile_OpcodeStackEffect`.
+Added *jump* parameter to :func:`dis.stack_effect`.