Handle converting StopIteration to RuntimeError in bytecode.

Author: markshannon

Include/internal/pycore_opcode.h

@@ -425,6 +425,9 @@ def _write_atomic(path, data, mode=0o666):
 #     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.13 will start with 3550
 
 # MAGIC must change whenever the bytecode emitted by the compiler may no
@@ -436,7 +439,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3510).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3514).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Include/opcode.h

@@ -190,7 +190,7 @@ def pseudo_op(name, op, real_ops):
 def_op('DELETE_DEREF', 139)
 hasfree.append(139)
 jrel_op('JUMP_BACKWARD', 140)    # Number of words to skip (backwards)
-
+def_op('LOAD_ERROR', 141)
 def_op('CALL_FUNCTION_EX', 142)  # Flags
 
 def_op('EXTENDED_ARG', 144)

Lib/importlib/_bootstrap_external.py

@@ -543,8 +543,16 @@ async def _asyncwith(c):
         >> COPY                     3
            POP_EXCEPT
            RERAISE                  1
+        >> LOAD_ERROR               1
+           CHECK_EXC_MATCH
+           POP_JUMP_IF_FALSE        8 (to 140)
+           LOAD_ERROR               2
+           LOAD_CONST               3 ('coroutine raised StopIteration')
+           CALL                     0
+           RAISE_VARARGS            2
+        >> RERAISE                  1
 ExceptionTable:
-6 rows
+12 rows
 """ % (_asyncwith.__code__.co_firstlineno,
        _asyncwith.__code__.co_firstlineno + 1,
        _asyncwith.__code__.co_firstlineno + 2,
@@ -1256,6 +1264,7 @@ def f(c=c):
 Constants:
    0: None
    1: <code object f at (.*), file "(.*)", line (.*)>
+   2: 'generator raised StopIteration'
 Variable names:
    0: a
    1: b
@@ -1362,6 +1371,7 @@ async def async_def():
 Constants:
    0: None
    1: 1
+   2: 'coroutine raised StopIteration'
 Names:
    0: b
    1: c

Lib/opcode.py

@@ -1439,7 +1439,7 @@ def bar(cls):
             check(bar, size('PP'))
         # generator
         def get_gen(): yield 1
-        check(get_gen(), size('P2P4P4c7P2ic??P'))
+        check(get_gen(), size('P2P4P4c7P2ic??4P'))
         # iterator
         check(iter('abc'), size('lP'))
         # callable-iterator

Lib/test/test_dis.py

@@ -346,7 +346,7 @@ def make_tracer():
         return Tracer()
 
     def compare_events(self, line_offset, events, expected_events):
-        events = [(l - line_offset, e) for (l, e) in events]
+        events = [(l - line_offset if l is not None else None, e) for (l, e) in events]
         if events != expected_events:
             self.fail(
                 "events did not match expectation:\n" +

Lib/test/test_sys.py

@@ -2175,6 +2175,26 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             goto exception_unwind;
         }
 
+        TARGET(LOAD_ERROR) {
+            PyObject *value;
+            switch(oparg) {
+                case 0:
+                    value = PyExc_AssertionError;
+                    break;
+                case 1:
+                    value = PyExc_StopIteration;
+                    break;
+                case 2:
+                    value = PyExc_RuntimeError;
+                    break;
+                default:
+                    Py_UNREACHABLE();
+            }
+            Py_INCREF(value);
+            PUSH(value);
+            DISPATCH();
+        }
+
         TARGET(LOAD_ASSERTION_ERROR) {
             PyObject *value = PyExc_AssertionError;
             Py_INCREF(value);

Lib/test/test_sys_settrace.py

@@ -1254,6 +1254,7 @@ stack_effect(int opcode, int oparg, int jump)
             return (oparg & FVS_MASK) == FVS_HAVE_SPEC ? -1 : 0;
         case LOAD_METHOD:
             return 1;
+        case LOAD_ERROR:
         case LOAD_ASSERTION_ERROR:
             return 1;
         case LIST_TO_TUPLE:
@@ -2567,6 +2568,42 @@ compiler_check_debug_args(struct compiler *c, arguments_ty args)
     return 1;
 }
 
+static int
+coroutine_stopiteration_handler(struct compiler *c, jump_target_label *handler)
+{
+    ADDOP_LOAD_CONST(c, NO_LOCATION, Py_None);
+    ADDOP(c, NO_LOCATION, RETURN_VALUE);
+    if (cfg_builder_use_label(CFG_BUILDER(c), *handler) < 0) {
+        return 0;
+    }
+    jump_target_label other = cfg_new_label(CFG_BUILDER(c));
+    if (!IS_LABEL(other)) {
+        return 0;
+    }
+    USE_LABEL(c, *handler);
+    ADDOP_I(c, NO_LOCATION, LOAD_ERROR, 1); // StopIteration
+    ADDOP(c, NO_LOCATION, CHECK_EXC_MATCH);
+    ADDOP_JUMP(c, NO_LOCATION, POP_JUMP_IF_FALSE, other);
+    ADDOP_I(c, NO_LOCATION, LOAD_ERROR, 2); // RuntimeError
+    const char *msg = c->u->u_ste->ste_coroutine ?
+        (c->u->u_ste->ste_generator ?
+            "async generator raised StopIteration" :
+            "coroutine raised StopIteration"
+        ) :
+        "generator raised StopIteration";
+    PyObject *message = _PyUnicode_FromASCII(msg, strlen(msg));
+    if (message == NULL) {
+        return 0;
+    }
+    ADDOP_LOAD_CONST_NEW(c, NO_LOCATION, message);
+    ADDOP_I(c, NO_LOCATION, CALL, 0);
+    ADDOP_I(c, NO_LOCATION, RAISE_VARARGS, 2);
+
+    USE_LABEL(c, other);
+    ADDOP_I(c, NO_LOCATION, RERAISE, 1);
+    return 1;
+}
+
 static int
 compiler_function(struct compiler *c, stmt_ty s, int is_async)
 {
@@ -2632,6 +2669,17 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
         return 0;
     }
 
+    bool is_coro = (c->u->u_ste->ste_coroutine || c->u->u_ste->ste_generator);
+    jump_target_label handler;
+    if (is_coro) {
+        handler = cfg_new_label(CFG_BUILDER(c));
+        if (!IS_LABEL(handler)) {
+            compiler_exit_scope(c);
+            return 0;
+        }
+        ADDOP_JUMP(c, NO_LOCATION, SETUP_CLEANUP, handler);
+    }
+
     /* if not -OO mode, add docstring */
     if (c->c_optimize < 2) {
         docstring = _PyAST_GetDocString(body);
@@ -2647,6 +2695,12 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
     for (i = docstring ? 1 : 0; i < asdl_seq_LEN(body); i++) {
         VISIT_IN_SCOPE(c, stmt, (stmt_ty)asdl_seq_GET(body, i));
     }
+    if (is_coro) {
+        if (!coroutine_stopiteration_handler(c, &handler)) {
+            compiler_exit_scope(c);
+            return 0;
+        }
+    }
     co = assemble(c, 1);
     qualname = c->u->u_qualname;
     Py_INCREF(qualname);