Add custom instruction for converting StopIteration into RuntimeError.

python · markshannon · Nov 3, 2022 · Nov 1, 2022 · Nov 2, 2022 · Nov 2, 2022
commit 194722704577eab0849d17cd6c3d3b2157d81fcf
diff --git a/Include/internal/pycore_opcode.h b/Include/internal/pycore_opcode.h
diff --git a/Include/opcode.h b/Include/opcode.h
@@ -424,9 +424,7 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.12a1 3508 (Add CLEANUP_THROW)
 #     Python 3.12a1 3509 (Conditional jumps only jump forward)
 #     Python 3.12a1 3510 (FOR_ITER leaves iterator on the stack)
-
-
-#     Python 3.12a1 3514 (Add LOAD_ERROR instruction)
+#     Python 3.12a1 3511 (Add STOPITERATION_ERROR instruction)
 
 #     Python 3.13 will start with 3550
 

diff --git a/Lib/opcode.py b/Lib/opcode.py
@@ -111,6 +111,8 @@ def pseudo_op(name, op, real_ops):
 def_op('STORE_SUBSCR', 60)
 def_op('DELETE_SUBSCR', 61)
 
+def_op('STOPITERATION_ERROR', 63)
+
 def_op('GET_ITER', 68)
 def_op('GET_YIELD_FROM_ITER', 69)
 def_op('PRINT_EXPR', 70)

diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -543,16 +543,8 @@ async def _asyncwith(c):
         >> COPY                     3
            POP_EXCEPT
            RERAISE                  1
-        >> LOAD_ERROR               1
-           CHECK_EXC_MATCH
-           POP_JUMP_IF_FALSE       10 (to 144)
-           PUSH_EXC_INFO
-           LOAD_ERROR               2
-           LOAD_CONST               3 ('coroutine raised StopIteration')
-           CALL                     0
-           SWAP                     2
-           RAISE_VARARGS            2
-        >> RERAISE                  1
+        >> STOPITERATION_ERROR
+           RERAISE                  1
 ExceptionTable:
 12 rows
 """ % (_asyncwith.__code__.co_firstlineno,
@@ -712,18 +704,6 @@ def foo(x):
            JUMP_BACKWARD            9 (to 8)
         >> END_FOR
            RETURN_VALUE
-        >> LOAD_ERROR               1
-           CHECK_EXC_MATCH
-           POP_JUMP_IF_FALSE       10 (to 56)
-           PUSH_EXC_INFO
-           LOAD_ERROR               2
-           LOAD_CONST               1 ('generator raised StopIteration')
-           CALL                     0
-           SWAP                     2
-           RAISE_VARARGS            2
-        >> RERAISE                  1
-ExceptionTable:
-1 row
 """ % (dis_nested_1,
        __file__,
        _h.__code__.co_firstlineno + 3,
@@ -1278,7 +1258,6 @@ def f(c=c):
 Constants:
    0: None
    1: <code object f at (.*), file "(.*)", line (.*)>
-   2: 'generator raised StopIteration'
 Variable names:
    0: a
    1: b
@@ -1385,7 +1364,6 @@ async def async_def():
 Constants:
    0: None
    1: 1
-   2: 'coroutine raised StopIteration'
 Names:
    0: b
    1: c

diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py
@@ -1439,7 +1439,7 @@ def bar(cls):
             check(bar, size('PP'))
         # generator
         def get_gen(): yield 1
-        check(get_gen(), size('P2P4P4c7P2ic??5P'))
+        check(get_gen(), size('P2P4P4c7P2ic??2P'))
         # iterator
         check(iter('abc'), size('lP'))
         # callable-iterator

diff --git a/Python/ceval.c b/Python/ceval.c
@@ -2198,6 +2198,48 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             DISPATCH();
         }
 
+        TARGET(STOPITERATION_ERROR) {
+            assert(frame->owner == FRAME_OWNED_BY_GENERATOR);
+            PyObject *exc = TOP();
+            assert(PyExceptionInstance_Check(exc));
+            const char *msg = NULL;
+            if (PyErr_GivenExceptionMatches(exc, PyExc_StopIteration)) {
+                msg = "generator raised StopIteration";
+                if (frame->f_code->co_flags & CO_ASYNC_GENERATOR) {
+                    msg = "async generator raised StopIteration";
+                }
+                else if (frame->f_code->co_flags & CO_COROUTINE) {
+                    msg = "coroutine raised StopIteration";
+                }
+            }
+            else if ((frame->f_code->co_flags & CO_ASYNC_GENERATOR) &&
+                    PyErr_GivenExceptionMatches(exc, PyExc_StopAsyncIteration))
+            {
+                /* code in `gen` raised a StopAsyncIteration error:
+                raise a RuntimeError.
+                */
+                msg = "async generator raised StopAsyncIteration";
+            }
+            if (msg != NULL) {
+                PyObject *message = _PyUnicode_FromASCII(msg, strlen(msg));
+                if (message == NULL) {
+                    goto error;
+                }
+                PyObject *error = PyObject_CallOneArg(PyExc_RuntimeError, message);
+                if (error == NULL) {
+                    Py_DECREF(message);
+                    goto error;
+                }
+                assert(PyExceptionInstance_Check(error));
+                SET_TOP(error);
+                PyException_SetCause(error, exc);
+                Py_INCREF(exc);
+                PyException_SetContext(error, exc);
+                Py_DECREF(message);
+            }
+            DISPATCH();
+        }
+
         TARGET(LOAD_ASSERTION_ERROR) {
             PyObject *value = PyExc_AssertionError;
             Py_INCREF(value);