Use a virtual try/except to handle close()/throw()

brandtbucher · Aug 19, 2022 · Apr 1, 2022 · May 6, 2022 · May 6, 2022 · May 6, 2022
commit c05c2af6db2cdc3bb89d9ee3fd479e87ba5869ba
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
@@ -344,6 +344,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(top)
         STRUCT_FOR_ID(truncate)
         STRUCT_FOR_ID(unraisablehook)
+        STRUCT_FOR_ID(value)
         STRUCT_FOR_ID(values)
         STRUCT_FOR_ID(version)
         STRUCT_FOR_ID(warnings)

diff --git a/Include/internal/pycore_runtime_init.h b/Include/internal/pycore_runtime_init.h
@@ -967,6 +967,7 @@ extern "C" {
                 INIT_ID(top), \
                 INIT_ID(truncate), \
                 INIT_ID(unraisablehook), \
+                INIT_ID(value), \
                 INIT_ID(values), \
                 INIT_ID(version), \
                 INIT_ID(warnings), \

diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -1023,9 +1023,11 @@ async def async_def():
 Constants:
    0: None
    1: 1
+   2: <class 'StopIteration'>
 Names:
-   0: b
-   1: c
+   0: value
+   1: b
+   2: c
 Variable names:
    0: a
    1: d"""

diff --git a/Objects/exceptions.c b/Objects/exceptions.c
@@ -585,17 +585,23 @@ StopIteration_traverse(PyStopIterationObject *self, visitproc visit, void *arg)
     return BaseException_traverse((PyBaseExceptionObject *)self, visit, arg);
 }
 
-ComplexExtendsException(
-    PyExc_Exception,       /* base */
-    StopIteration,         /* name */
-    StopIteration,         /* prefix for *_init, etc */
-    0,                     /* new */
-    0,                     /* methods */
-    StopIteration_members, /* members */
-    0,                     /* getset */
-    0,                     /* str */
-    "Signal the end from iterator.__next__()."
-);
+// Don't use ComplexExtendsException for this, since deepfreeze doesn't work if
+// the type is static:
+PyTypeObject _PyExc_StopIteration = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    .tp_name = "StopIteration",
+    .tp_basicsize = sizeof(PyStopIterationObject),
+    .tp_dealloc = (destructor)StopIteration_dealloc,
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
+    .tp_doc = PyDoc_STR("Signal the end from iterator.__next__()."),
+    .tp_traverse = (traverseproc)StopIteration_traverse,
+    .tp_clear = (inquiry)StopIteration_clear,
+    .tp_members = StopIteration_members,
+    .tp_base = &_PyExc_Exception,
+    .tp_dictoffset = offsetof(PyStopIterationObject, dict),
+    .tp_init = (initproc)StopIteration_init,
+};
+PyObject *PyExc_StopIteration = (PyObject *)&_PyExc_StopIteration;
 
 
 /*

@@ -478,28 +478,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(frame->f_lasti >= 0);
-            _Py_CODEUNIT *code = _PyCode_CODE(gen->gi_code);
-            /* Backup to SEND */
-            frame->f_lasti--;
-            assert(_Py_OPCODE(code[frame->f_lasti]) == SEND);
-            int jump = _Py_OPARG(code[frame->f_lasti]);
-            frame->f_lasti += jump;
-            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;
     }

@@ -2580,6 +2580,9 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
         }
 
         TARGET(YIELD_VALUE) {
+            // NOTE: It's important that YIELD_VALUE never raises an exception!
+            // The compiler treats *anything* raised here as part of a close()
+            // or throw() call.
             assert(frame->is_entry);
             PyObject *retval = POP();
             _PyFrame_GetGenerator(frame)->gi_frame_state = FRAME_SUSPENDED;

@@ -1875,19 +1875,37 @@ compiler_call_exit_with_nones(struct compiler *c) {
 static int
 compiler_add_yield_from(struct compiler *c, int await)
 {
-    basicblock *start, *resume, *exit;
-    start = compiler_new_block(c);
-    resume = compiler_new_block(c);
-    exit = compiler_new_block(c);
-    if (start == NULL || resume == NULL || exit == NULL) {
-        return 0;
-    }
+    basicblock *start, *resume, *stopiter, *error, *exit;
+    RETURN_IF_FALSE(start = compiler_new_block(c));
+    RETURN_IF_FALSE(resume = compiler_new_block(c));
+    RETURN_IF_FALSE(stopiter = compiler_new_block(c));
+    RETURN_IF_FALSE(error = compiler_new_block(c));
+    RETURN_IF_FALSE(exit = compiler_new_block(c));
+
     compiler_use_next_block(c, start);
     ADDOP_JUMP(c, SEND, exit);
     compiler_use_next_block(c, resume);
+    // Set up a virtual try/except to handle StopIteration raised during a
+    // close() or throw():
+    ADDOP_JUMP(c, SETUP_FINALLY, stopiter);
+    RETURN_IF_FALSE(compiler_push_fblock(c, TRY_EXCEPT, resume, NULL, NULL));
+    // The only way YIELD_VALUE can raise is if close() or throw() raises:
     ADDOP(c, YIELD_VALUE);
+    compiler_pop_fblock(c, TRY_EXCEPT, resume);
+    ADDOP_NOLINE(c, POP_BLOCK);
     ADDOP_I(c, RESUME, await ? 3 : 2);
     ADDOP_JUMP(c, JUMP_NO_INTERRUPT, start);
+    compiler_use_next_block(c, stopiter);
+    ADDOP_LOAD_CONST(c, PyExc_StopIteration);  // StopIteration is marshallable!
+    ADDOP_JUMP(c, JUMP_IF_NOT_EXC_MATCH, error);
+    // StopIteration was raised. Push the return value and continue execution:
+    ADDOP_NAME(c, LOAD_ATTR, &_Py_ID(value), names);
+    ADDOP_I(c, SWAP, 3);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP_JUMP(c, JUMP_FORWARD, exit);
+    compiler_use_next_block(c, error);
+    ADDOP_I(c, RERAISE, 0);
     compiler_use_next_block(c, exit);
     return 1;
 }

diff --git a/Tools/scripts/deepfreeze.py b/Tools/scripts/deepfreeze.py
@@ -117,6 +117,8 @@ def __init__(self, file: TextIO) -> None:
         self.write('#include "internal/pycore_code.h"')
         self.write('#include "internal/pycore_long.h"')
         self.write("")
+        self.write("extern PyTypeObject _PyExc_StopIteration;")
+        self.write("")
 
     @contextlib.contextmanager
     def indent(self) -> None:
@@ -404,6 +406,8 @@ def generate(self, name: str, obj: object) -> str:
             return "Py_Ellipsis"
         elif obj is None:
             return "Py_None"
+        elif obj is StopIteration:
+            return "(PyObject *)&_PyExc_StopIteration"
         else:
             raise TypeError(
                 f"Cannot generate code for {type(obj).__name__} object")

diff --git a/Tools/scripts/umarshal.py b/Tools/scripts/umarshal.py
@@ -298,6 +298,8 @@ def R_REF(obj: Any) -> Any:
             retval = self.refs[n]
             assert retval is not None
             return retval
+        elif type == Type.STOPITER:
+            return StopIteration
         else:
             breakpoint()
             raise AssertionError(f"Unknown type {type} {chr(type)!r}")
-Original file line number
+Diff line change
@@ Expand Up / @@ -1023,9 +1023,11 @@ async def async_def(): @@
     Constants:
 : None
 : 1
+: <class 'StopIteration'>
     Names:
-: b
-: c
+: value
+: b
+: c
     Variable names:
 : a
 : d"""
@@ Expand Down @@