GH-90997: Wrap yield from/await in a virtual try/except StopIteration (GH-96010)

Author: brandtbucher

Doc/library/dis.rst

@@ -567,6 +567,17 @@ the original TOS1.
     .. versionchanged:: 3.11
        Exception representation on the stack now consist of one, not three, items.
 
+
+.. opcode:: CLEANUP_THROW
+
+   Handles an exception raised during a :meth:`~generator.throw` or
+   :meth:`~generator.close` call through the current frame.  If TOS is an
+   instance of :exc:`StopIteration`, pop three values from the stack and push
+   its ``value`` member.  Otherwise, re-raise TOS.
+
+   .. versionadded:: 3.12
+
+
 .. opcode:: BEFORE_ASYNC_WITH
 
    Resolves ``__aenter__`` and ``__aexit__`` from the object on top of the
@@ -1344,10 +1355,14 @@ iterations of the loop.
     .. versionadded:: 3.11
 
 
-.. opcode:: SEND
+.. opcode:: SEND (delta)
+
+    Equivalent to ``TOS = TOS1.send(TOS)``. Used in ``yield from`` and ``await``
+    statements.
 
-    Sends ``None`` to the sub-generator of this generator.
-    Used in ``yield from`` and ``await`` statements.
+    If the call raises :exc:`StopIteration`, pop both items, push the
+    exception's ``value`` attribute, and increment the bytecode counter by
+    *delta*.
 
     .. versionadded:: 3.11
 

Include/internal/pycore_opcode.h

@@ -411,10 +411,10 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.12a1 3505 (Specialization/Cache for FOR_ITER)
 #     Python 3.12a1 3506 (Add BINARY_SLICE and STORE_SLICE instructions)
 #     Python 3.12a1 3507 (Set lineno of module's RESUME to 0)
+#     Python 3.12a1 3508 (Add CLEANUP_THROW)
 
 #     Python 3.13 will start with 3550
 
-#
 # MAGIC must change whenever the bytecode emitted by the compiler may no
 # longer be understood by older implementations of the eval loop (usually
 # due to the addition of new opcodes).
@@ -424,7 +424,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 = (3507).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3508).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Include/opcode.h

@@ -104,6 +104,7 @@ def pseudo_op(name, op, real_ops):
 def_op('BEFORE_ASYNC_WITH', 52)
 def_op('BEFORE_WITH', 53)
 def_op('END_ASYNC_FOR', 54)
+def_op('CLEANUP_THROW', 55)
 
 def_op('STORE_SUBSCR', 60)
 def_op('DELETE_SUBSCR', 61)

Lib/importlib/_bootstrap_external.py

@@ -507,26 +507,31 @@ async def _asyncwith(c):
            LOAD_CONST               0 (None)
            RETURN_VALUE
 
-%3d     >> PUSH_EXC_INFO
+%3d     >> CLEANUP_THROW
+           JUMP_BACKWARD           24 (to 22)
+        >> CLEANUP_THROW
+           JUMP_BACKWARD            9 (to 56)
+        >> PUSH_EXC_INFO
            WITH_EXCEPT_START
            GET_AWAITABLE            2
            LOAD_CONST               0 (None)
-        >> SEND                     3 (to 82)
+        >> SEND                     4 (to 92)
            YIELD_VALUE              6
            RESUME                   3
-           JUMP_BACKWARD_NO_INTERRUPT     4 (to 74)
-        >> POP_JUMP_FORWARD_IF_TRUE     1 (to 86)
+           JUMP_BACKWARD_NO_INTERRUPT     4 (to 82)
+        >> CLEANUP_THROW
+        >> POP_JUMP_FORWARD_IF_TRUE     1 (to 96)
            RERAISE                  2
         >> POP_TOP
            POP_EXCEPT
            POP_TOP
            POP_TOP
-           JUMP_BACKWARD           19 (to 58)
+           JUMP_BACKWARD           24 (to 58)
         >> COPY                     3
            POP_EXCEPT
            RERAISE                  1
 ExceptionTable:
-2 rows
+6 rows
 """ % (_asyncwith.__code__.co_firstlineno,
        _asyncwith.__code__.co_firstlineno + 1,
        _asyncwith.__code__.co_firstlineno + 2,

Lib/opcode.py

@@ -0,0 +1,3 @@
+Compile virtual :keyword:`try`/:keyword:`except` blocks to handle exceptions
+raised during :meth:`~generator.close` or :meth:`~generator.throw` calls
+through a suspended frame.

Lib/test/test_dis.py

@@ -485,26 +485,7 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
         }
         Py_DECREF(yf);
         if (!ret) {
-            PyObject *val;
-            /* Pop subiterator from stack */
-            assert(gen->gi_frame_state < FRAME_CLEARED);
-            ret = _PyFrame_StackPop((_PyInterpreterFrame *)gen->gi_iframe);
-            assert(ret == yf);
-            Py_DECREF(ret);
-            // XXX: Performing this jump ourselves is awkward and problematic.
-            // See https://github.com/python/cpython/pull/31968.
-            /* Termination repetition of SEND loop */
-            assert(_PyInterpreterFrame_LASTI(frame) >= 0);
-            /* Backup to SEND */
-            assert(_Py_OPCODE(frame->prev_instr[-1]) == SEND);
-            int jump = _Py_OPARG(frame->prev_instr[-1]);
-            frame->prev_instr += jump - 1;
-            if (_PyGen_FetchStopIterationValue(&val) == 0) {
-                ret = gen_send(gen, val);
-                Py_DECREF(val);
-            } else {
-                ret = gen_send_ex(gen, Py_None, 1, 0);
-            }
+            ret = gen_send_ex(gen, Py_None, 1, 0);
         }
         return ret;
     }