From e8bb05394164e7735f7a9de80a046953606a38eb Mon Sep 17 00:00:00 2001
From: Jacob Bower <1978924+jbower-fb@users.noreply.github.com>
Date: Thu, 21 Nov 2024 15:37:49 -0800
Subject: [PATCH] gh-126091: Always link generator frames when propagating a
 thrown-in exception through a yield-from chain (#126092)

Always link generator frames when propagating a thrown-in exception through a yield-from chain.
---
 Lib/test/test_generators.py                   | 22 ++++++++++++++++++-
 ...-11-07-21-48-23.gh-issue-126091.ETaRGE.rst |  2 ++
 Objects/genobject.c                           | 18 ++++++++++-----
 3 files changed, 35 insertions(+), 7 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core_and_Builtins/2024-11-07-21-48-23.gh-issue-126091.ETaRGE.rst

diff --git a/Lib/test/test_generators.py b/Lib/test/test_generators.py
index bf2cb116072..2ea6dba12ef 100644
--- a/Lib/test/test_generators.py
+++ b/Lib/test/test_generators.py
@@ -758,7 +758,8 @@ class GeneratorStackTraceTest(unittest.TestCase):
         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 @@ class GeneratorStackTraceTest(unittest.TestCase):
 
         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
new file mode 100644
index 00000000000..08118ff1af6
--- /dev/null
+++ 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
index 19c2c4e3331..e87f199c250 100644
--- 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);