Add @asynccontextmanager and get tests passing again

Author: njsmith

async_generator/__init__.py

@@ -6,9 +6,9 @@
     isasyncgen,
     isasyncgenfunction,
 )
-from ._util import aclosing
+from ._util import aclosing, asynccontextmanager
 
 __all__ = [
     "async_generator", "yield_", "yield_from_", "aclosing", "isasyncgen",
-    "isasyncgenfunction"
+    "isasyncgenfunction", "asynccontextmanager",
 ]

async_generator/_tests/test_async_generator.py

@@ -7,7 +7,7 @@
 from functools import wraps
 import gc
 
-from . import (
+from .. import (
     async_generator,
     yield_,
     yield_from_,
@@ -717,7 +717,7 @@ async def f():
 # Finicky tests to check that the overly clever ctype stuff has plausible
 # refcounting
 
-from . import impl
+from .. import _impl
 
 
 @pytest.mark.skipif(not hasattr(sys, "getrefcount"), reason="CPython only")
@@ -728,12 +728,12 @@ def test_refcnt():
     print(sys.getrefcount(x))
     print(sys.getrefcount(x))
     base_count = sys.getrefcount(x)
-    l = [impl._wrap(x) for _ in range(100)]
+    l = [_impl._wrap(x) for _ in range(100)]
     print(sys.getrefcount(x))
     print(sys.getrefcount(x))
     print(sys.getrefcount(x))
     assert sys.getrefcount(x) >= base_count + 100
-    l2 = [impl._unwrap(box) for box in l]
+    l2 = [_impl._unwrap(box) for box in l]
     assert sys.getrefcount(x) >= base_count + 200
     print(sys.getrefcount(x))
     print(sys.getrefcount(x))

async_generator/_tests/test_util.py

@@ -1,6 +1,6 @@
 import pytest
 
-from . import aclosing, async_generator, yield_
+from .. import aclosing, async_generator, yield_, asynccontextmanager
 
 
 @async_generator
@@ -34,3 +34,158 @@ async def test_aclosing():
     except ValueError:
         pass
     assert closed_slot[0]
+
+
+@pytest.mark.asyncio
+async def test_contextmanager_do_not_unchain_non_stopiteration_exceptions():
+    @asynccontextmanager
+    @async_generator
+    async def manager_issue29692():
+        try:
+            await yield_()
+        except Exception as exc:
+            raise RuntimeError('issue29692:Chained') from exc
+
+    with pytest.raises(RuntimeError) as excinfo:
+        async with manager_issue29692():
+            raise ZeroDivisionError
+    assert excinfo.value.args[0] == 'issue29692:Chained'
+    assert isinstance(excinfo.value.__cause__, ZeroDivisionError)
+
+    # This is a little funky because of implementation details in
+    # async_generator It can all go away once we stop supporting Python3.5
+    with pytest.raises(RuntimeError) as excinfo:
+        async with manager_issue29692():
+            exc = StopIteration('issue29692:Unchained')
+            raise exc
+    assert excinfo.value.args[0] == 'issue29692:Chained'
+    cause = excinfo.value.__cause__
+    assert cause.args[0] == 'generator raised StopIteration'
+    assert cause.__cause__ is exc
+
+    with pytest.raises(StopAsyncIteration) as excinfo:
+        async with manager_issue29692():
+            raise StopAsyncIteration('issue29692:Unchained')
+    assert excinfo.value.args[0] == 'issue29692:Unchained'
+    assert excinfo.value.__cause__ is None
+
+    @asynccontextmanager
+    @async_generator
+    async def noop_async_context_manager():
+        await yield_()
+
+    with pytest.raises(StopIteration):
+        async with noop_async_context_manager():
+            raise StopIteration
+
+
+# Native async generators are only available from Python 3.6 and onwards
+nativeasyncgenerators = True
+try:
+    exec(
+        """
+@asynccontextmanager
+async def manager_issue29692_2():
+    try:
+        yield
+    except Exception as exc:
+        raise RuntimeError('issue29692:Chained') from exc
+"""
+    )
+except SyntaxError:
+    nativeasyncgenerators = False
+
+
+@pytest.mark.skipif(
+    not nativeasyncgenerators,
+    reason="Python < 3.6 doesn't have native async generators"
+)
+@pytest.mark.asyncio
+async def test_native_contextmanager_do_not_unchain_non_stopiteration_exceptions(
+):
+
+    with pytest.raises(RuntimeError) as excinfo:
+        async with manager_issue29692_2():
+            raise ZeroDivisionError
+    assert excinfo.value.args[0] == 'issue29692:Chained'
+    assert isinstance(excinfo.value.__cause__, ZeroDivisionError)
+
+    for cls in [StopIteration, StopAsyncIteration]:
+        with pytest.raises(cls) as excinfo:
+            async with manager_issue29692_2():
+                raise cls('issue29692:Unchained')
+        assert excinfo.value.args[0] == 'issue29692:Unchained'
+        assert excinfo.value.__cause__ is None
+
+
+@pytest.mark.asyncio
+async def test_asynccontextmanager_exception_passthrough():
+    # This was the cause of annoying coverage flapping, see gh-140
+    @asynccontextmanager
+    @async_generator
+    async def noop_async_context_manager():
+        await yield_()
+
+    for exc_type in [StopAsyncIteration, RuntimeError, ValueError]:
+        with pytest.raises(exc_type):
+            async with noop_async_context_manager():
+                raise exc_type
+
+
+@pytest.mark.asyncio
+async def test_asynccontextmanager_catches_exception():
+    @asynccontextmanager
+    @async_generator
+    async def catch_it():
+        with pytest.raises(ValueError):
+            await yield_()
+
+    async with catch_it():
+        raise ValueError
+
+
+@pytest.mark.asyncio
+async def test_asynccontextmanager_no_yield():
+    @asynccontextmanager
+    @async_generator
+    async def yeehaw():
+        pass
+
+    with pytest.raises(RuntimeError) as excinfo:
+        async with yeehaw():
+            assert False  # pragma: no cover
+
+    assert "didn't yield" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_asynccontextmanager_too_many_yields():
+    @asynccontextmanager
+    @async_generator
+    async def doubleyield():
+        try:
+            await yield_()
+        except Exception:
+            pass
+        await yield_()
+
+    with pytest.raises(RuntimeError) as excinfo:
+        async with doubleyield():
+            pass
+
+    assert "didn't stop" in str(excinfo.value)
+
+    with pytest.raises(RuntimeError) as excinfo:
+        async with doubleyield():
+            raise ValueError
+
+    assert "didn't stop after athrow" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_asynccontextmanager_requires_asyncgenfunction():
+    with pytest.raises(TypeError):
+
+        @asynccontextmanager
+        def syncgen():  # pragma: no cover
+            yield

async_generator/_util.py

@@ -1,3 +1,7 @@
+import sys
+from functools import wraps
+from ._impl import isasyncgenfunction
+
 class aclosing:
     def __init__(self, aiter):
         self._aiter = aiter
@@ -7,3 +11,97 @@ async def __aenter__(self):
 
     async def __aexit__(self, *args):
         await self._aiter.aclose()
+
+
+# Very much derived from the one in contextlib, by copy/pasting and then
+# asyncifying everything. (Also I dropped the obscure support for using
+# context managers as function decorators. It could be re-added; I just
+# couldn't be bothered.)
+# So this is a derivative work licensed under the PSF License, which requires
+# the following notice:
+#
+# Copyright © 2001-2017 Python Software Foundation; All Rights Reserved
+class _AsyncGeneratorContextManager:
+    def __init__(self, func, args, kwds):
+        self._func_name = func.__name__
+        self._agen = func(*args, **kwds).__aiter__()
+
+    async def __aenter__(self):
+        if sys.version_info < (3, 5, 2):
+            self._agen = await self._agen
+        try:
+            return await self._agen.asend(None)
+        except StopAsyncIteration:
+            raise RuntimeError("async generator didn't yield") from None
+
+    async def __aexit__(self, type, value, traceback):
+        if type is None:
+            try:
+                await self._agen.asend(None)
+            except StopAsyncIteration:
+                return False
+            else:
+                raise RuntimeError("async generator didn't stop")
+        else:
+            # It used to be possible to have type != None, value == None:
+            #    https://bugs.python.org/issue1705170
+            # but AFAICT this can't happen anymore.
+            assert value is not None
+            try:
+                await self._agen.athrow(type, value, traceback)
+                raise RuntimeError(
+                    "async 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. (issue27112)
+                if exc is value:
+                    return False
+                # Likewise, avoid suppressing if a StopIteration exception
+                # was passed to throw() and later wrapped into a RuntimeError
+                # (see PEP 479).
+                if (isinstance(value, (StopIteration, StopAsyncIteration))
+                        and exc.__cause__ is value):
+                    return False
+                raise
+            except:
+                # 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 sys.exc_info()[1] is value:
+                    return False
+                raise
+
+    def __enter__(self):
+        raise RuntimeError(
+            "use 'async with {func_name}(...)', not 'with {func_name}(...)'".
+            format(func_name=self._func_name)
+        )
+
+    def __exit__(self):  # pragma: no cover
+        assert False, """Never called, but should be defined"""
+
+
+def asynccontextmanager(func):
+    """Like @contextmanager, but async."""
+    if not isasyncgenfunction(func):
+        raise TypeError(
+            "must be an async generator (native or from async_generator; "
+            "if using @async_generator then @acontextmanager must be on top."
+        )
+
+    @wraps(func)
+    def helper(*args, **kwds):
+        return _AsyncGeneratorContextManager(func, args, kwds)
+
+    # A hint for sphinxcontrib-trio:
+    helper.__returns_acontextmanager__ = True
+    return helper