From dcba50f18b9885322031b969fbea20aa95ef33b2 Mon Sep 17 00:00:00 2001
From: Jacob Bower <jbower@meta.com>
Date: Sat, 7 Dec 2024 15:57:17 -0800
Subject: [PATCH] Backport gh-126091: Always link generator frames when
 propagating a thrown-in exception through a yield-from chain

Reviewed By: alexmalyshev

Differential Revision: D66556446

fbshipit-source-id: 022f87d78d439efb1595e2c638cf91b976ab9ab1
---
 Objects/genobject.c | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/Objects/genobject.c b/Objects/genobject.c
index e632f43e92a..899c84326d7 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -641,18 +641,15 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
                 return gen_send_ex(gen, Py_None, 1, 0);
             goto throw_here;
         }
+
+        PyThreadState *tstate = _PyThreadState_GET();
+        PyFrameObject *f = tstate->frame;
+        _PyShadowFrame *sf = tstate->shadow_frame;
+        _PyShadowFrame_PtrKind old_ptr_kind = _PyShadowFrame_GetPtrKind(sf);
         if (PyGen_CheckExact(yf) || PyCoro_CheckExact(yf)) {
             /* `yf` is a generator or a coroutine. */
-            PyThreadState *tstate = _PyThreadState_GET();
-            PyFrameObject *f = tstate->frame;
-            _PyShadowFrame *sf = tstate->shadow_frame;
-            _PyShadowFrame_PtrKind old_ptr_kind = _PyShadowFrame_GetPtrKind(sf);
-
-            /* 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. */
             tstate->frame = gen->gi_frame;
             tstate->shadow_frame = &gen->gi_shadow_frame;
             /* Close the generator that we are currently iterating with
@@ -678,9 +675,17 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
                 Py_DECREF(yf);
                 goto throw_here;
             }
+            tstate->frame = gen->gi_frame;
+            tstate->shadow_frame = &gen->gi_shadow_frame;
             CI_WITH_GEN_MARKED_AS_THROWING(gen, {
                 ret = PyObject_CallFunctionObjArgs(meth, typ, val, tb, NULL);
             });
+            if (old_ptr_kind != PYSF_PYFRAME && _PyShadowFrame_GetPtrKind(sf) == PYSF_PYFRAME) {
+              /* Frame was materialized while throwing into yf */
+              f = _PyShadowFrame_GetPyFrame(sf);
+            }
+            tstate->frame = f;
+            tstate->shadow_frame = sf;
             Py_DECREF(meth);
         }
         Py_DECREF(yf);