diff options
| -rw-r--r-- | Lib/contextlib.py | 8 | ||||
| -rw-r--r-- | Lib/test/test_contextlib.py | 34 | ||||
| -rw-r--r-- | Lib/test/test_contextlib_async.py | 35 | ||||
| -rw-r--r-- | Misc/NEWS.d/next/Library/2021-07-12-10-32-48.bpo-44594.eEa5zi.rst | 3 |
4 files changed, 76 insertions, 4 deletions
diff --git a/Lib/contextlib.py b/Lib/contextlib.py index 1e0a39f..c63a849 100644 --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -540,10 +540,10 @@ class ExitStack(_BaseExitStack, AbstractContextManager): # Context may not be correct, so find the end of the chain while 1: exc_context = new_exc.__context__ - if exc_context is old_exc: + if exc_context is None or exc_context is old_exc: # Context is already set correctly (see issue 20317) return - if exc_context is None or exc_context is frame_exc: + if exc_context is frame_exc: break new_exc = exc_context # Change the end of the chain to point to the exception @@ -674,10 +674,10 @@ class AsyncExitStack(_BaseExitStack, AbstractAsyncContextManager): # Context may not be correct, so find the end of the chain while 1: exc_context = new_exc.__context__ - if exc_context is old_exc: + if exc_context is None or exc_context is old_exc: # Context is already set correctly (see issue 20317) return - if exc_context is None or exc_context is frame_exc: + if exc_context is frame_exc: break new_exc = exc_context # Change the end of the chain to point to the exception diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py index 5a08065..fbaae2d 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -780,6 +780,40 @@ class TestBaseExitStack: self.assertIsInstance(inner_exc, ValueError) self.assertIsInstance(inner_exc.__context__, ZeroDivisionError) + def test_exit_exception_explicit_none_context(self): + # Ensure ExitStack chaining matches actual nested `with` statements + # regarding explicit __context__ = None. + + class MyException(Exception): + pass + + @contextmanager + def my_cm(): + try: + yield + except BaseException: + exc = MyException() + try: + raise exc + finally: + exc.__context__ = None + + @contextmanager + def my_cm_with_exit_stack(): + with self.exit_stack() as stack: + stack.enter_context(my_cm()) + yield stack + + for cm in (my_cm, my_cm_with_exit_stack): + with self.subTest(): + try: + with cm(): + raise IndexError() + except MyException as exc: + self.assertIsNone(exc.__context__) + else: + self.fail("Expected IndexError, but no exception was raised") + def test_exit_exception_non_suppressing(self): # http://bugs.python.org/issue19092 def raise_exc(exc): diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py index 603162e..127d750 100644 --- a/Lib/test/test_contextlib_async.py +++ b/Lib/test/test_contextlib_async.py @@ -552,6 +552,41 @@ class TestAsyncExitStack(TestBaseExitStack, unittest.TestCase): self.assertIsInstance(inner_exc, ValueError) self.assertIsInstance(inner_exc.__context__, ZeroDivisionError) + @_async_test + async def test_async_exit_exception_explicit_none_context(self): + # Ensure AsyncExitStack chaining matches actual nested `with` statements + # regarding explicit __context__ = None. + + class MyException(Exception): + pass + + @asynccontextmanager + async def my_cm(): + try: + yield + except BaseException: + exc = MyException() + try: + raise exc + finally: + exc.__context__ = None + + @asynccontextmanager + async def my_cm_with_exit_stack(): + async with self.exit_stack() as stack: + await stack.enter_async_context(my_cm()) + yield stack + + for cm in (my_cm, my_cm_with_exit_stack): + with self.subTest(): + try: + async with cm(): + raise IndexError() + except MyException as exc: + self.assertIsNone(exc.__context__) + else: + self.fail("Expected IndexError, but no exception was raised") + class TestAsyncNullcontext(unittest.TestCase): @_async_test diff --git a/Misc/NEWS.d/next/Library/2021-07-12-10-32-48.bpo-44594.eEa5zi.rst b/Misc/NEWS.d/next/Library/2021-07-12-10-32-48.bpo-44594.eEa5zi.rst new file mode 100644 index 0000000..a2bfd8f --- /dev/null +++ b/Misc/NEWS.d/next/Library/2021-07-12-10-32-48.bpo-44594.eEa5zi.rst @@ -0,0 +1,3 @@ +Fix an edge case of :class:`ExitStack` and :class:`AsyncExitStack` exception +chaining. They will now match ``with`` block behavior when ``__context__`` is +explicitly set to ``None`` when the exception is in flight. |