diff options
| -rw-r--r-- | Lib/contextlib.py | 3 | ||||
| -rw-r--r-- | Lib/test/test_contextlib.py | 47 |
2 files changed, 50 insertions, 0 deletions
diff --git a/Lib/contextlib.py b/Lib/contextlib.py index ef8f8c9..1ff8cdf 100644 --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -105,6 +105,9 @@ class _GeneratorContextManager(_GeneratorContextManagerBase, 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 index 96ceaa7..2a44404 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -8,6 +8,7 @@ import threading import unittest from contextlib import * # Tests __all__ from test import support +import weakref class TestAbstractContextManager(unittest.TestCase): @@ -219,6 +220,52 @@ def woohoo(): 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): |