Close #13585: add contextlib.ExitStack to replace the ill-fated contextlib.nested API
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index 6e38305..8bed88e 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -370,6 +370,129 @@
self.assertEqual(state, [1, 'something else', 999])
+class TestExitStack(unittest.TestCase):
+
+ def test_no_resources(self):
+ with ExitStack():
+ pass
+
+ def test_callback(self):
+ expected = [
+ ((), {}),
+ ((1,), {}),
+ ((1,2), {}),
+ ((), dict(example=1)),
+ ((1,), dict(example=1)),
+ ((1,2), dict(example=1)),
+ ]
+ result = []
+ def _exit(*args, **kwds):
+ """Test metadata propagation"""
+ result.append((args, kwds))
+ with ExitStack() as stack:
+ for args, kwds in reversed(expected):
+ if args and kwds:
+ f = stack.callback(_exit, *args, **kwds)
+ elif args:
+ f = stack.callback(_exit, *args)
+ elif kwds:
+ f = stack.callback(_exit, **kwds)
+ else:
+ f = stack.callback(_exit)
+ self.assertIs(f, _exit)
+ for wrapper in stack._exit_callbacks:
+ self.assertIs(wrapper.__wrapped__, _exit)
+ self.assertNotEqual(wrapper.__name__, _exit.__name__)
+ self.assertIsNone(wrapper.__doc__, _exit.__doc__)
+ self.assertEqual(result, expected)
+
+ def test_push(self):
+ exc_raised = ZeroDivisionError
+ def _expect_exc(exc_type, exc, exc_tb):
+ self.assertIs(exc_type, exc_raised)
+ def _suppress_exc(*exc_details):
+ return True
+ def _expect_ok(exc_type, exc, exc_tb):
+ self.assertIsNone(exc_type)
+ self.assertIsNone(exc)
+ self.assertIsNone(exc_tb)
+ class ExitCM(object):
+ def __init__(self, check_exc):
+ self.check_exc = check_exc
+ def __enter__(self):
+ self.fail("Should not be called!")
+ def __exit__(self, *exc_details):
+ self.check_exc(*exc_details)
+ with ExitStack() as stack:
+ stack.push(_expect_ok)
+ self.assertIs(stack._exit_callbacks[-1], _expect_ok)
+ cm = ExitCM(_expect_ok)
+ stack.push(cm)
+ self.assertIs(stack._exit_callbacks[-1].__self__, cm)
+ stack.push(_suppress_exc)
+ self.assertIs(stack._exit_callbacks[-1], _suppress_exc)
+ cm = ExitCM(_expect_exc)
+ stack.push(cm)
+ self.assertIs(stack._exit_callbacks[-1].__self__, cm)
+ stack.push(_expect_exc)
+ self.assertIs(stack._exit_callbacks[-1], _expect_exc)
+ stack.push(_expect_exc)
+ self.assertIs(stack._exit_callbacks[-1], _expect_exc)
+ 1/0
+
+ def test_enter_context(self):
+ class TestCM(object):
+ def __enter__(self):
+ result.append(1)
+ def __exit__(self, *exc_details):
+ result.append(3)
+
+ result = []
+ cm = TestCM()
+ with ExitStack() as stack:
+ @stack.callback # Registered first => cleaned up last
+ def _exit():
+ result.append(4)
+ self.assertIsNotNone(_exit)
+ stack.enter_context(cm)
+ self.assertIs(stack._exit_callbacks[-1].__self__, cm)
+ result.append(2)
+ self.assertEqual(result, [1, 2, 3, 4])
+
+ def test_close(self):
+ result = []
+ with ExitStack() as stack:
+ @stack.callback
+ def _exit():
+ result.append(1)
+ self.assertIsNotNone(_exit)
+ stack.close()
+ result.append(2)
+ self.assertEqual(result, [1, 2])
+
+ def test_pop_all(self):
+ result = []
+ with ExitStack() as stack:
+ @stack.callback
+ def _exit():
+ result.append(3)
+ self.assertIsNotNone(_exit)
+ new_stack = stack.pop_all()
+ result.append(1)
+ result.append(2)
+ new_stack.close()
+ self.assertEqual(result, [1, 2, 3])
+
+ def test_instance_bypass(self):
+ class Example(object): pass
+ cm = Example()
+ cm.__exit__ = object()
+ stack = ExitStack()
+ self.assertRaises(AttributeError, stack.enter_context, cm)
+ stack.push(cm)
+ self.assertIs(stack._exit_callbacks[-1], cm)
+
+
# This is needed to make the test actually run under regrtest.py!
def test_main():
support.run_unittest(__name__)