diff --git a/Lib/contextlib.py b/Lib/contextlib.py index 004d1037b78a476..63dd1dedcf7f584 100644 --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -531,10 +531,10 @@ def _fix_exception_context(new_exc, old_exc): # 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 @@ -671,10 +671,10 @@ def _fix_exception_context(new_exc, old_exc): # 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 9c27866cd661cb1..825197f26824b0d 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -794,6 +794,40 @@ def suppress_exc(*exc_details): 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 7904abff7d1aa57..004c6deef4bf4fd 100644 --- a/Lib/test/test_contextlib_async.py +++ b/Lib/test/test_contextlib_async.py @@ -556,6 +556,41 @@ async def suppress_exc(*exc_details): 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") + @_async_test async def test_instance_bypass_async(self): class Example(object): pass 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 000000000000000..a2bfd8ff5b51bc0 --- /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.