Skip to content

Commit dd0e087

Browse files
teckincoghlan
authored andcommitted
bpo-30306: release arguments of contextmanager (GH-1500)
The arguments to a generator function which is declared as a contextmanager are stored inside the context manager, and thus are kept alive, even when it is used as a regular context manager, and not as a function decorator (where it needs the original arguments to recreate the generator on each call). This is a possible unnecessary memory leak, so this changes contextmanager.__enter__ to release the saved arguments, as that method being called means that particular CM instance isn't going to need to recreate the underlying generator. Patch by Martin Teichmann.
1 parent c4b1248 commit dd0e087

File tree

2 files changed

+50
-0
lines changed

2 files changed

+50
-0
lines changed

Lib/contextlib.py

+3
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,9 @@ def _recreate_cm(self):
105105
return self.__class__(self.func, self.args, self.kwds)
106106

107107
def __enter__(self):
108+
# do not keep args and kwds alive unnecessarily
109+
# they are only needed for recreation, which is not possible anymore
110+
del self.args, self.kwds, self.func
108111
try:
109112
return next(self.gen)
110113
except StopIteration:

Lib/test/test_contextlib.py

+47
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import unittest
99
from contextlib import * # Tests __all__
1010
from test import support
11+
import weakref
1112

1213

1314
class TestAbstractContextManager(unittest.TestCase):
@@ -219,6 +220,52 @@ def woohoo(self, func, args, kwds):
219220
with woohoo(self=11, func=22, args=33, kwds=44) as target:
220221
self.assertEqual(target, (11, 22, 33, 44))
221222

223+
def test_nokeepref(self):
224+
class A:
225+
pass
226+
227+
@contextmanager
228+
def woohoo(a, b):
229+
a = weakref.ref(a)
230+
b = weakref.ref(b)
231+
self.assertIsNone(a())
232+
self.assertIsNone(b())
233+
yield
234+
235+
with woohoo(A(), b=A()):
236+
pass
237+
238+
def test_param_errors(self):
239+
@contextmanager
240+
def woohoo(a, *, b):
241+
yield
242+
243+
with self.assertRaises(TypeError):
244+
woohoo()
245+
with self.assertRaises(TypeError):
246+
woohoo(3, 5)
247+
with self.assertRaises(TypeError):
248+
woohoo(b=3)
249+
250+
def test_recursive(self):
251+
depth = 0
252+
@contextmanager
253+
def woohoo():
254+
nonlocal depth
255+
before = depth
256+
depth += 1
257+
yield
258+
depth -= 1
259+
self.assertEqual(depth, before)
260+
261+
@woohoo()
262+
def recursive():
263+
if depth < 10:
264+
recursive()
265+
266+
recursive()
267+
self.assertEqual(depth, 0)
268+
222269

223270
class ClosingTestCase(unittest.TestCase):
224271

0 commit comments

Comments
 (0)