pythongh-118272: Clear generator frame's locals when the generator is closed (python#118277)

Co-authored-by: Thomas Grainger <[email protected]>

Author: 2 people

Include/internal/pycore_frame.h

@@ -227,6 +227,9 @@ _PyFrame_GetFrameObject(_PyInterpreterFrame *frame)
     return _PyFrame_MakeAndSetFrameObject(frame);
 }
 
+void
+_PyFrame_ClearLocals(_PyInterpreterFrame *frame);
+
 /* Clears all references in the frame.
  * If take is non-zero, then the _PyInterpreterFrame frame
  * may be transferred to the frame object it references

Lib/test/test_generators.py

@@ -532,6 +532,26 @@ def f():
         with self.assertRaises(RuntimeError):
             gen.close()
 
+    def test_close_releases_frame_locals(self):
+        # See gh-118272
+
+        class Foo:
+            pass
+
+        f = Foo()
+        f_wr = weakref.ref(f)
+
+        def genfn():
+            a = f
+            yield
+
+        g = genfn()
+        next(g)
+        del f
+        g.close()
+        support.gc_collect()
+        self.assertIsNone(f_wr())
+
 
 class GeneratorThrowTest(unittest.TestCase):
 

Misc/NEWS.d/next/Core and Builtins/2024-04-25-12-55-47.gh-issue-118272.5ptjk_.rst

@@ -0,0 +1,2 @@
+Fix bug where ``generator.close`` does not free the generator frame's
+locals.

Objects/genobject.c

@@ -380,6 +380,7 @@ gen_close(PyGenObject *gen, PyObject *args)
             // RESUME after YIELD_VALUE and exception depth is 1
             assert((oparg & RESUME_OPARG_LOCATION_MASK) != RESUME_AT_FUNC_START);
             gen->gi_frame_state = FRAME_COMPLETED;
+            _PyFrame_ClearLocals((_PyInterpreterFrame *)gen->gi_iframe);
             Py_RETURN_NONE;
         }
     }

Python/frame.c

@@ -94,6 +94,17 @@ take_ownership(PyFrameObject *f, _PyInterpreterFrame *frame)
     }
 }
 
+void
+_PyFrame_ClearLocals(_PyInterpreterFrame *frame)
+{
+    assert(frame->stacktop >= 0);
+    for (int i = 0; i < frame->stacktop; i++) {
+        Py_XDECREF(frame->localsplus[i]);
+    }
+    frame->stacktop = 0;
+    Py_CLEAR(frame->f_locals);
+}
+
 void
 _PyFrame_ClearExceptCode(_PyInterpreterFrame *frame)
 {
@@ -114,11 +125,7 @@ _PyFrame_ClearExceptCode(_PyInterpreterFrame *frame)
         }
         Py_DECREF(f);
     }
-    assert(frame->stacktop >= 0);
-    for (int i = 0; i < frame->stacktop; i++) {
-        Py_XDECREF(frame->localsplus[i]);
-    }
-    Py_XDECREF(frame->f_locals);
+    _PyFrame_ClearLocals(frame);
     Py_DECREF(frame->f_funcobj);
 }