gh-126091: Always link generator frames when propagating a thrown-in exception through a yield-from chain (#126092)

jbower-fb · web-flow · commit e8bb05394164 · 2024-11-21T17:37:49.000-06:00
Always link generator frames when propagating a thrown-in exception through a yield-from chain.
diff --git a/Lib/test/test_generators.py b/Lib/test/test_generators.py
@@ -758,7 +758,8 @@ def check_stack_names(self, frame, expected):
         while frame:
             name = frame.f_code.co_name
             # Stop checking frames when we get to our test helper.
-            if name.startswith('check_') or name.startswith('call_'):
+            if (name.startswith('check_') or name.startswith('call_')
+                    or name.startswith('test')):
                 break
 
             names.append(name)
@@ -799,6 +800,25 @@ def call_throw(gen):
 
         self.check_yield_from_example(call_throw)
 
+    def test_throw_with_yield_from_custom_generator(self):
+
+        class CustomGen:
+            def __init__(self, test):
+                self.test = test
+            def throw(self, *args):
+                self.test.check_stack_names(sys._getframe(), ['throw', 'g'])
+            def __iter__(self):
+                return self
+            def __next__(self):
+                return 42
+
+        def g(target):
+            yield from target
+
+        gen = g(CustomGen(self))
+        gen.send(None)
+        gen.throw(RuntimeError)
+
 
 class YieldFromTests(unittest.TestCase):
     def test_generator_gi_yieldfrom(self):
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-11-07-21-48-23.gh-issue-126091.ETaRGE.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-11-07-21-48-23.gh-issue-126091.ETaRGE.rst
@@ -0,0 +1,2 @@
+Ensure stack traces are complete when throwing into a generator chain that
+ends in a custom generator.
diff --git a/Objects/genobject.c b/Objects/genobject.c
@@ -471,14 +471,14 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
                 return gen_send_ex(gen, Py_None, 1, 0);
             goto throw_here;
         }
+        PyThreadState *tstate = _PyThreadState_GET();
+        assert(tstate != NULL);
         if (PyGen_CheckExact(yf) || PyCoro_CheckExact(yf)) {
             /* `yf` is a generator or a coroutine. */
-            PyThreadState *tstate = _PyThreadState_GET();
-            /* Since we are fast-tracking things by skipping the eval loop,
-               we need to update the current frame so the stack trace
-               will be reported correctly to the user. */
-            /* XXX We should probably be updating the current frame
-               somewhere in ceval.c. */
+
+            /* Link frame into the stack to enable complete backtraces. */
+            /* XXX We should probably be updating the current frame somewhere in
+               ceval.c. */
             _PyInterpreterFrame *prev = tstate->current_frame;
             frame->previous = prev;
             tstate->current_frame = frame;
@@ -502,10 +502,16 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
                 Py_DECREF(yf);
                 goto throw_here;
             }
+
+            _PyInterpreterFrame *prev = tstate->current_frame;
+            frame->previous = prev;
+            tstate->current_frame = frame;
             PyFrameState state = gen->gi_frame_state;
             gen->gi_frame_state = FRAME_EXECUTING;
             ret = PyObject_CallFunctionObjArgs(meth, typ, val, tb, NULL);
             gen->gi_frame_state = state;
+            tstate->current_frame = prev;
+            frame->previous = NULL;
             Py_DECREF(meth);
         }
         Py_DECREF(yf);

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+Ensure stack traces are complete when throwing into a generator chain that`
	`2`	`+ends in a custom generator.`