From 7f1c330da31c54e028dceaf3610877914c2a4497 Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Tue, 20 Jul 2021 19:15:07 +0100 Subject: [PATCH] bpo-44566: resolve differences between asynccontextmanager and contextmanager (#27024) --- Lib/contextlib.py | 104 ++++++++++-------- Lib/test/test_contextlib.py | 23 ++-- Lib/test/test_contextlib_async.py | 13 ++- .../2021-07-05-18-13-25.bpo-44566.o51Bd1.rst | 1 + 4 files changed, 85 insertions(+), 56 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2021-07-05-18-13-25.bpo-44566.o51Bd1.rst diff --git a/Lib/contextlib.py b/Lib/contextlib.py index 004d1037b78a47..8343d7e5196713 100644 --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -113,18 +113,20 @@ def __init__(self, func, args, kwds): # for the class instead. # See http://bugs.python.org/issue19404 for more details. - -class _GeneratorContextManager(_GeneratorContextManagerBase, - AbstractContextManager, - ContextDecorator): - """Helper for @contextmanager decorator.""" - def _recreate_cm(self): - # _GCM instances are one-shot context managers, so the + # _GCMB instances are one-shot context managers, so the # CM must be recreated each time a decorated function is # called return self.__class__(self.func, self.args, self.kwds) + +class _GeneratorContextManager( + _GeneratorContextManagerBase, + AbstractContextManager, + ContextDecorator, +): + """Helper for @contextmanager decorator.""" + def __enter__(self): # do not keep args and kwds alive unnecessarily # they are only needed for recreation, which is not possible anymore @@ -134,8 +136,8 @@ def __enter__(self): except StopIteration: raise RuntimeError("generator didn't yield") from None - def __exit__(self, type, value, traceback): - if type is None: + def __exit__(self, typ, value, traceback): + if typ is None: try: next(self.gen) except StopIteration: @@ -146,9 +148,9 @@ def __exit__(self, type, value, traceback): if value is None: # Need to force instantiation so we can reliably # tell if we get the same exception back - value = type() + value = typ() try: - self.gen.throw(type, value, traceback) + self.gen.throw(typ, value, traceback) except StopIteration as exc: # Suppress StopIteration *unless* it's the same exception that # was passed to throw(). This prevents a StopIteration @@ -158,81 +160,93 @@ def __exit__(self, type, value, traceback): # Don't re-raise the passed in exception. (issue27122) if exc is value: return False - # Likewise, avoid suppressing if a StopIteration exception + # Avoid suppressing if a StopIteration exception # was passed to throw() and later wrapped into a RuntimeError - # (see PEP 479). - if type is StopIteration and exc.__cause__ is value: + # (see PEP 479 for sync generators; async generators also + # have this behavior). But do this only if the exception wrapped + # by the RuntimeError is actually Stop(Async)Iteration (see + # issue29692). + if ( + isinstance(value, StopIteration) + and exc.__cause__ is value + ): return False raise - except: + except BaseException as exc: # only re-raise if it's *not* the exception that was # passed to throw(), because __exit__() must not raise # an exception unless __exit__() itself failed. But throw() # has to raise the exception to signal propagation, so this # fixes the impedance mismatch between the throw() protocol # and the __exit__() protocol. - # - # This cannot use 'except BaseException as exc' (as in the - # async implementation) to maintain compatibility with - # Python 2, where old-style class exceptions are not caught - # by 'except BaseException'. - if sys.exc_info()[1] is value: - return False - raise + if exc is not value: + raise + return False raise RuntimeError("generator didn't stop after throw()") - -class _AsyncGeneratorContextManager(_GeneratorContextManagerBase, - AbstractAsyncContextManager, - AsyncContextDecorator): - """Helper for @asynccontextmanager.""" - - def _recreate_cm(self): - # _AGCM instances are one-shot context managers, so the - # ACM must be recreated each time a decorated function is - # called - return self.__class__(self.func, self.args, self.kwds) +class _AsyncGeneratorContextManager( + _GeneratorContextManagerBase, + AbstractAsyncContextManager, + AsyncContextDecorator, +): + """Helper for @asynccontextmanager decorator.""" async def __aenter__(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 await self.gen.__anext__() + return await anext(self.gen) except StopAsyncIteration: raise RuntimeError("generator didn't yield") from None async def __aexit__(self, typ, value, traceback): if typ is None: try: - await self.gen.__anext__() + await anext(self.gen) except StopAsyncIteration: - return + return False else: raise RuntimeError("generator didn't stop") else: if value is None: + # Need to force instantiation so we can reliably + # tell if we get the same exception back value = typ() - # See _GeneratorContextManager.__exit__ for comments on subtleties - # in this implementation try: await self.gen.athrow(typ, value, traceback) - raise RuntimeError("generator didn't stop after athrow()") except StopAsyncIteration as exc: + # Suppress StopIteration *unless* it's the same exception that + # was passed to throw(). This prevents a StopIteration + # raised inside the "with" statement from being suppressed. return exc is not value except RuntimeError as exc: + # Don't re-raise the passed in exception. (issue27122) if exc is value: return False - # Avoid suppressing if a StopIteration exception - # was passed to throw() and later wrapped into a RuntimeError + # Avoid suppressing if a Stop(Async)Iteration exception + # was passed to athrow() and later wrapped into a RuntimeError # (see PEP 479 for sync generators; async generators also # have this behavior). But do this only if the exception wrapped # by the RuntimeError is actually Stop(Async)Iteration (see # issue29692). - if isinstance(value, (StopIteration, StopAsyncIteration)): - if exc.__cause__ is value: - return False + if ( + isinstance(value, (StopIteration, StopAsyncIteration)) + and exc.__cause__ is value + ): + return False raise except BaseException as exc: + # only re-raise if it's *not* the exception that was + # passed to throw(), because __exit__() must not raise + # an exception unless __exit__() itself failed. But throw() + # has to raise the exception to signal propagation, so this + # fixes the impedance mismatch between the throw() protocol + # and the __exit__() protocol. if exc is not value: raise + return False + raise RuntimeError("generator didn't stop after athrow()") def contextmanager(func): diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py index 9c27866cd661cb..04720d949fee5b 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -126,19 +126,22 @@ def woohoo(): self.assertEqual(state, [1, 42, 999]) def test_contextmanager_except_stopiter(self): - stop_exc = StopIteration('spam') @contextmanager def woohoo(): yield - try: - with self.assertWarnsRegex(DeprecationWarning, - "StopIteration"): - with woohoo(): - raise stop_exc - except Exception as ex: - self.assertIs(ex, stop_exc) - else: - self.fail('StopIteration was suppressed') + + class StopIterationSubclass(StopIteration): + pass + + for stop_exc in (StopIteration('spam'), StopIterationSubclass('spam')): + with self.subTest(type=type(stop_exc)): + try: + with woohoo(): + raise stop_exc + except Exception as ex: + self.assertIs(ex, stop_exc) + else: + self.fail(f'{stop_exc} was suppressed') def test_contextmanager_except_pep479(self): code = """\ diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py index 7904abff7d1aa5..6a218f911569c9 100644 --- a/Lib/test/test_contextlib_async.py +++ b/Lib/test/test_contextlib_async.py @@ -209,7 +209,18 @@ async def test_contextmanager_except_stopiter(self): async def woohoo(): yield - for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')): + class StopIterationSubclass(StopIteration): + pass + + class StopAsyncIterationSubclass(StopAsyncIteration): + pass + + for stop_exc in ( + StopIteration('spam'), + StopAsyncIteration('ham'), + StopIterationSubclass('spam'), + StopAsyncIterationSubclass('spam') + ): with self.subTest(type=type(stop_exc)): try: async with woohoo(): diff --git a/Misc/NEWS.d/next/Library/2021-07-05-18-13-25.bpo-44566.o51Bd1.rst b/Misc/NEWS.d/next/Library/2021-07-05-18-13-25.bpo-44566.o51Bd1.rst new file mode 100644 index 00000000000000..3b00a1b715feef --- /dev/null +++ b/Misc/NEWS.d/next/Library/2021-07-05-18-13-25.bpo-44566.o51Bd1.rst @@ -0,0 +1 @@ +handle StopIteration subclass raised from @contextlib.contextmanager generator \ No newline at end of file