bpo-30306: release arguments of contextmanager (GH-1500)

tecki · ncoghlan · commit dd0e087edc8f · 2018-01-28T14:17:46.000+10:00
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.
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
@@ -105,6 +105,9 @@ def _recreate_cm(self):
         return self.__class__(self.func, self.args, self.kwds)
 
     def __enter__(self):
+        # do not keep args and kwds alive unnecessarily
+        # they are only needed for recreation, which is not possible anymore
+        del self.args, self.kwds, self.func
         try:
             return next(self.gen)
         except StopIteration:
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
@@ -8,6 +8,7 @@
 import unittest
 from contextlib import *  # Tests __all__
 from test import support
+import weakref
 
 
 class TestAbstractContextManager(unittest.TestCase):
@@ -219,6 +220,52 @@ def woohoo(self, func, args, kwds):
         with woohoo(self=11, func=22, args=33, kwds=44) as target:
             self.assertEqual(target, (11, 22, 33, 44))
 
+    def test_nokeepref(self):
+        class A:
+            pass
+
+        @contextmanager
+        def woohoo(a, b):
+            a = weakref.ref(a)
+            b = weakref.ref(b)
+            self.assertIsNone(a())
+            self.assertIsNone(b())
+            yield
+
+        with woohoo(A(), b=A()):
+            pass
+
+    def test_param_errors(self):
+        @contextmanager
+        def woohoo(a, *, b):
+            yield
+
+        with self.assertRaises(TypeError):
+            woohoo()
+        with self.assertRaises(TypeError):
+            woohoo(3, 5)
+        with self.assertRaises(TypeError):
+            woohoo(b=3)
+
+    def test_recursive(self):
+        depth = 0
+        @contextmanager
+        def woohoo():
+            nonlocal depth
+            before = depth
+            depth += 1
+            yield
+            depth -= 1
+            self.assertEqual(depth, before)
+
+        @woohoo()
+        def recursive():
+            if depth < 10:
+                recursive()
+
+        recursive()
+        self.assertEqual(depth, 0)
+
 
 class ClosingTestCase(unittest.TestCase):
 
