Skip to content

Commit

Permalink
pythongh-118272: Clear generator frame's locals when the generator is…
Browse files Browse the repository at this point in the history
… closed
  • Loading branch information
iritkatriel committed Apr 25, 2024
1 parent f180b31 commit 66519f6
Show file tree
Hide file tree
Showing 5 changed files with 38 additions and 5 deletions.
3 changes: 3 additions & 0 deletions Include/internal/pycore_frame.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
20 changes: 20 additions & 0 deletions Lib/test/test_generators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
gc.collect()
self.assertIsNone(f_wr())


class GeneratorThrowTest(unittest.TestCase):

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix bug where ``generator.close`` does not free the generator frame's
locals.
1 change: 1 addition & 0 deletions Objects/genobject.c
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
Expand Down
17 changes: 12 additions & 5 deletions Python/frame.c
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand All @@ -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);
}

Expand Down

0 comments on commit 66519f6

Please sign in to comment.