Close #13585: add contextlib.ExitStack to replace the ill-fated contextlib.nested API

1 files changed, 123 insertions, 0 deletions

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 @@ class TestContextDecorator(unittest.TestCase):
         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__)