bpo-36851: Clean the frame stack if the execution ends with a return and the stack is not empty (GH-13191)

1 files changed, 38 insertions, 0 deletions

diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
index e49121e..9bf290d 100644
--- a/Lib/test/test_code.py
+++ b/Lib/test/test_code.py
@@ -130,6 +130,7 @@ import sys
 import threading
 import unittest
 import weakref
+import opcode
 try:
     import ctypes
 except ImportError:
@@ -379,6 +380,43 @@ if check_impl_detail(cpython=True) and ctypes is not None:
             tt.join()
             self.assertEqual(LAST_FREED, 500)
 
+        @cpython_only
+        def test_clean_stack_on_return(self):
+
+            def f(x):
+                return x
+
+            code = f.__code__
+            ct = type(f.__code__)
+
+            # Insert an extra LOAD_FAST, this duplicates the value of
+            # 'x' in the stack, leaking it if the frame is not properly
+            # cleaned up upon exit.
+
+            bytecode = list(code.co_code)
+            bytecode.insert(-2, opcode.opmap['LOAD_FAST'])
+            bytecode.insert(-2, 0)
+
+            c = ct(code.co_argcount, code.co_posonlyargcount,
+                   code.co_kwonlyargcount, code.co_nlocals, code.co_stacksize+1,
+                   code.co_flags, bytes(bytecode),
+                   code.co_consts, code.co_names, code.co_varnames,
+                   code.co_filename, code.co_name, code.co_firstlineno,
+                   code.co_lnotab, code.co_freevars, code.co_cellvars)
+            new_function = type(f)(c, f.__globals__, 'nf', f.__defaults__, f.__closure__)
+
+            class Var:
+                pass
+            the_object = Var()
+            var = weakref.ref(the_object)
+
+            new_function(the_object)
+
+            # Check if the_object is leaked
+            del the_object
+            assert var() is None
+
+
 def test_main(verbose=None):
     from test import test_code
     run_doctest(test_code, verbose)