Repository: python-trio/trio

Issue:
Nurseries cannot propagate StopAsyncIteration exceptions
I'm trying to implement an async version of `zip` using the following code, but I'm getting an error when one of the iterators passed to `async_zip` raises a `StopAsyncIteration` before the others are finished. Normally I'd expect that I can catch it using the multierror or a try/except, however, instead I'm getting a `RuntimeError`.

I don't understand why this is happening and it's not clear if this is a bug in my code or in trio.

```py
import trio

class async_zip(object):

    def __init__(self, *largs):
        self.nexts = [obj.__anext__ for obj in largs]

    async def _accumulate(self, f, items, i):
        items[i] = await f()

    def __aiter__(self):
        return self

    async def __anext__(self):
        nexts = self.nexts
        items = [None, ] * len(nexts)
        got_stop = False

        def handle(exc):
            nonlocal got_stop
            if isinstance(exc, StopAsyncIteration):
                got_stop = True
                return None
            else:
                return exc

        with trio.MultiError.catch(handle):
            async with trio.open_nursery() as nursery:
                for i, f in enumerate(nexts):
                    nursery.start_soon(self._accumulate, f, items, i)

        if got_stop:
            raise StopAsyncIteration
        return items

# an async iterable
class it(object):
    def __init__(self, count):
        self.count = count
        self.val = 0
    def __aiter__(self):
        return self
    async def __anext__(self):
        await trio.sleep(1)
        val = self.val
        if val >= self.count:
            raise StopAsyncIteration
        self.val += 1 
        return val

# test the iterable
async def runner():
    async for val in it(4):
        print('got', val)

# let all iterables finish together
async def zipper():
    async for vals in async_zip(it(4), it(4), it(4), it(4)):
        print('got', vals)

# have one iterable finish early
async def zipper_short():
    async for vals in async_zip(it(4), it(4), it(4), it(2)):
        print('got', vals)

# run them
trio.run(runner)
trio.run(zipper)
trio.run(zipper_short)
```

This prints

```
got 0
got 1
got 2
got 3
got [0, 0, 0, 0]
got [1, 1, 1, 1]
got [2, 2, 2, 2]
got [3, 3, 3, 3]
got [0, 0, 0, 0]
got [1, 1, 1, 1]
Traceback (most recent call last):
  File "E:\Python\Python35-x64\lib\site-packages\async_generator\impl.py", line 199, in _invoke
    result = fn(*args)
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 320, in open_nursery
    await nursery._clean_up(pending_exc)
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 519, in _clean_up
    raise mexc
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 1502, in run_impl
    msg = task.coro.send(next_send)
  File "<ipython-input-2-e43ada780e94>", line 9, in _accumulate
    items[i] = await f()
  File "<ipython-input-2-e43ada780e94>", line 46, in __anext__
    raise StopAsyncIteration
StopAsyncIteration
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
  File "E:\Python\Python35-x64\lib\site-packages\IPython\core\interactiveshell.py", line 2881, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-2-e43ada780e94>", line 64, in <module>
    trio.run(zipper_short)
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 1397, in run
    return result.unwrap()
  File "E:\Python\Python35-x64\lib\contextlib.py", line 77, in __exit__
    self.gen.throw(type, value, traceback)
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_ki.py", line 194, in ki_manager
    yield
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 1397, in run
    return result.unwrap()
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_result.py", line 119, in unwrap
    raise self.error
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 1502, in run_impl
    msg = task.coro.send(next_send)
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 1085, in init
    return system_nursery._reap_and_unwrap(task)
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 447, in _reap_and_unwrap
    return task._result.unwrap()
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_result.py", line 119, in unwrap
    raise self.error
  File "E:\Python\Python35-x64\lib\site-packages\trio\_core\_run.py", line 1502, in run_impl
    msg = task.coro.send(next_send)
  File "<ipython-input-2-e43ada780e94>", line 59, in zipper_short
    async for vals in async_zip(it(4), it(4), it(4), it(2)):
  File "<ipython-input-2-e43ada780e94>", line 30, in __anext__
    nursery.start_soon(self._accumulate, f, items, i)
  File "E:\Python\Python35-x64\lib\site-packages\trio\_util.py", line 109, in __aexit__
    await self._agen.asend(None)
  File "E:\Python\Python35-x64\lib\site-packages\async_generator\impl.py", line 261, in asend
    return await self._do_it(self._it.send, value)
  File "E:\Python\Python35-x64\lib\site-packages\async_generator\impl.py", line 277, in _do_it
    return await ANextIter(self._it, start_fn, *args)
  File "E:\Python\Python35-x64\lib\site-packages\async_generator\impl.py", line 192, in send
    return self._invoke(self._it.send, value)
  File "E:\Python\Python35-x64\lib\site-packages\async_generator\impl.py", line 211, in _invoke
    ) from e
RuntimeError: async_generator raise StopAsyncIteration
```


Discussion:
Oh wow, this is a tricky bug you found. At first I thought it was #229, but it's actually different. (BTW, that concurrent zip is really cool! Even if it does seem to be uniquely suited to tickling this very weird bug :-).) Anyway, here's what's happening:

A popular way to implement context managers is by writing a generator, and then decorating it with `contextlib.contextmanager`. This same pattern can be used to write async context managers, by writing an async generator and then using an appropriately-implemented *async* context manager decorator. Trio has an internal implementation of this decorator (`trio._util.acontextmanager`), and that's how `trio.open_nursery` is implemented. So `open_nursery` looks roughly like:

```python
@acontextmanager
async def open_nursery():
    # ... do setup stuff
    try:
        yield nursery
    except BaseException as exc:
        nursery._record_exception(exc)
    await nursery._wait_for_all_tasks_to_exit()
    raise nursery._combine_exceptions()
```

Here, exactly one of the background tasks raised `StopAsyncIteration`, so the last line ends up being `raise that_stop_async_iteration_object`.

However, there's a problem: you're not allowed to raise `StopAsyncIteration` inside an async generator! According to PEP 479, if you try, then it gets replaced by a `RuntimeError`. This is actually a good thing, because the way generators work, a regular return is converted into a `StopAsyncIteration` (just like falling off the end of a generator function causes `StopIteration` to be raised), so if it wasn't converted into a `RuntimeError` then it would have ended up getting swallowed by the decorator and mis-interpreted as a bare `return`, which would have been even more confusing. But... either way, this isn't going to work.

Our `@acontextmanager` decorator has some special code to detect this problem when the exception was actually raised inside the body of the `async with` block, and fix things up again – which means that code like this works the way you'd expect:
```python
async with trio.open_nursery() as nursery:
    raise StopAsyncIteration
```
But the only reason this works is because the `@acontextmanager` catches the `StopAsyncIteration` first before throwing it into the `open_nursery` code, so it can recognize the special `RuntimeError` when it comes back out. In your case, the `StopAsyncIteration` is being raised by a background task and gets routed into the nursery via the nursery machinery, so `@acontextmanager` doesn't have any way to recognize it, and we can't use this special trick.

So this really depends on a bunch of arcane details of how async context managers, async generators, and the nursery implementation all interact. Neat!

But, uh... what do we do about it?

I think the only option is to give up on using `@acontextmanager` to implement `open_nursery`, and switch to implementing the async context manager protocol directly. This is a bit tricky, but I've actually [done most of the work already](https://github.com/python-trio/trio/blob/a351a2d148032dccf2b06636efbf4399bc0e24a4/trio/_core/_run.py#L321-L349) (see also #340). I guess this will be the thing that pushes us over into actually switching to that! (And this will also fix #229, so that's nice.)

While you're waiting, though, here's a workaround: the problem is specifically with having a function called by `nursery.start_soon` that raises `StopAsyncIteration`. If you handle the exception *inside* the function, instead of letting it escape, then that should work. So something like (untested):

```python
class async_zip(object):

    def __init__(self, *largs):
        # I added the __aiter__() call here; it's nothing to do with this bug, but you should have it :-)
        self.nexts = [obj.__aiter__().__anext__ for obj in largs]
        # We put the "are we done?" state on the object so _accumulate can see it
        self.done = False

    async def _accumulate(self, f, items, i):
        try:
            items[i] = await f()
        except StopAsyncIteration:
            self.done = True

    def __aiter__(self):
        return self

    async def __anext__(self):
        if self.done:
            # catch nasty users who try to re-iterate us after we said we were finished
            # (another advantage of having it on the object)
            raise StopAsyncIteration

        nexts = self.nexts
        items = [None, ] * len(nexts)

        async with trio.open_nursery() as nursery:
            for i, f in enumerate(nexts):
                nursery.start_soon(self._accumulate, f, items, i)

        if self.done:
            raise StopAsyncIteration
        # I also added the call to tuple() here because that's how zip() is supposed to work :-)
        return tuple(items)
```

Alternatively you could implement it as an async generator yourself (this is using the native async generator syntax, so requires Python 3.6+). Maybe this is clearer? I dunno:

```python
async def async_zip(*aiterables):
    aiterators = [aiterable.__aiter__() for aiterable in aiterables]
    done = False
    while True:
        items = [None] * len(aiterators)

        async def fill_in(i):
            try:
                items[i] = await aiterators[i].__anext__()
            except StopAsyncIteration:
                nonlocal done
                done = True

        async with trio.open_nursery() as nursery:
            for i in range(len(aiterators):
                nursery.start_soon(fill_in, i)

        if done:
            break
        yield tuple(items)
```
  
acontextmanager is also a missing piece of python stdlib.
Are you considering to upstream it?

It's already upstream, as contextlib.asynccontextmanager. (I had nothing to
do with this :-).) The problem is that this won't actually ship until
Python 3.7 comes out, which is ~6 months away, and longer than that to
become the minimum supported version. I'll probably move trio's version
into async_generator for those who need it sooner. (If someone wants to
make a PR for this then please do!)

On Jan 5, 2018 7:42 AM, "Imran Geriskovan" <notifications@github.com> wrote:

> acontextmanager is also a missing piece of python stdlib.
> Are you considering to upstream it?
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/python-trio/trio/issues/393#issuecomment-355585597>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AAlOaBrHiOscMgHjOW1HVNP64-NuVxY3ks5tHkLKgaJpZM4RUADX>
> .
>

Ahh, thanks for explaining that (and little code fixes). Your first work-around will work nicely (I'm still on 3.5.2).

Resolving patch:
diff --git a/trio/_core/_run.py b/trio/_core/_run.py
--- a/trio/_core/_run.py
+++ b/trio/_core/_run.py
@@ -16,9 +16,7 @@
 from sniffio import current_async_library_cvar
 
 import attr
-from async_generator import (
-    async_generator, yield_, asynccontextmanager, isasyncgen
-)
+from async_generator import isasyncgen
 from sortedcontainers import SortedDict
 from outcome import Error, Value
 
@@ -295,57 +293,55 @@ def started(self, value=None):
         self._old_nursery._check_nursery_closed()
 
 
-@asynccontextmanager
-@async_generator
-@enable_ki_protection
-async def open_nursery():
-    """Returns an async context manager which creates a new nursery.
+class NurseryManager:
+    """Nursery context manager.
 
-    This context manager's ``__aenter__`` method executes synchronously. Its
-    ``__aexit__`` method blocks until all child tasks have exited.
+    Note we explicitly avoid @asynccontextmanager and @async_generator
+    since they add a lot of extraneous stack frames to exceptions, as
+    well as cause problematic behavior with handling of StopIteration
+    and StopAsyncIteration.
 
     """
-    assert currently_ki_protected()
-    with open_cancel_scope() as scope:
-        nursery = Nursery(current_task(), scope)
-        nested_child_exc = None
-        try:
-            await yield_(nursery)
-        except BaseException as exc:
-            nested_child_exc = exc
+
+    @enable_ki_protection
+    async def __aenter__(self):
+        assert currently_ki_protected()
+        self._scope_manager = open_cancel_scope()
+        scope = self._scope_manager.__enter__()
+        self._nursery = Nursery(current_task(), scope)
+        return self._nursery
+
+    @enable_ki_protection
+    async def __aexit__(self, etype, exc, tb):
         assert currently_ki_protected()
-        await nursery._nested_child_finished(nested_child_exc)
-
-
-# I *think* this is equivalent to the above, and it gives *much* nicer
-# exception tracebacks... but I'm a little nervous about it because it's much
-# trickier code :-(
-#
-# class NurseryManager:
-#     @enable_ki_protection
-#     async def __aenter__(self):
-#         self._scope_manager = open_cancel_scope()
-#         scope = self._scope_manager.__enter__()
-#         self._parent_nursery = Nursery(current_task(), scope)
-#         return self._parent_nursery
-#
-#     @enable_ki_protection
-#     async def __aexit__(self, etype, exc, tb):
-#         try:
-#             await self._parent_nursery._clean_up(exc)
-#         except BaseException as new_exc:
-#             if not self._scope_manager.__exit__(
-#                     type(new_exc), new_exc, new_exc.__traceback__):
-#                 if exc is new_exc:
-#                     return False
-#                 else:
-#                     raise
-#         else:
-#             self._scope_manager.__exit__(None, None, None)
-#             return True
-#
-# def open_nursery():
-#     return NurseryManager()
+        try:
+            await self._nursery._nested_child_finished(exc)
+        except BaseException as new_exc:
+            try:
+                if self._scope_manager.__exit__(
+                    type(new_exc), new_exc, new_exc.__traceback__
+                ):
+                    return True
+            except BaseException as scope_manager_exc:
+                if scope_manager_exc == exc:
+                    return False
+                raise  # scope_manager_exc
+            raise  # new_exc
+        else:
+            self._scope_manager.__exit__(None, None, None)
+            return True
+
+    def __enter__(self):
+        raise RuntimeError(
+            "use 'async with open_nursery(...)', not 'with open_nursery(...)'"
+        )
+
+    def __exit__(self):  # pragma: no cover
+        assert False, """Never called, but should be defined"""
+
+
+def open_nursery():
+    return NurseryManager()
 
 
 class Nursery:


Test patch:
diff --git a/trio/_core/tests/test_run.py b/trio/_core/tests/test_run.py
--- a/trio/_core/tests/test_run.py
+++ b/trio/_core/tests/test_run.py
@@ -17,6 +17,7 @@
 from .tutil import check_sequence_matches, gc_collect_harder
 from ... import _core
 from ..._timeouts import sleep
+from ..._util import aiter_compat
 from ...testing import (
     wait_all_tasks_blocked,
     Sequencer,
@@ -1823,6 +1824,79 @@ async def start_sleep_then_crash(nursery):
         assert _core.current_time() - t0 == 7
 
 
+async def test_nursery_explicit_exception():
+    with pytest.raises(KeyError):
+        async with _core.open_nursery():
+            raise KeyError()
+
+
+async def test_nursery_stop_iteration():
+    async def fail():
+        raise ValueError
+
+    try:
+        async with _core.open_nursery() as nursery:
+            nursery.start_soon(fail)
+            raise StopIteration
+    except _core.MultiError as e:
+        assert tuple(map(type, e.exceptions)) == (StopIteration, ValueError)
+
+
+async def test_nursery_stop_async_iteration():
+    class it(object):
+        def __init__(self, count):
+            self.count = count
+            self.val = 0
+
+        async def __anext__(self):
+            await sleep(0)
+            val = self.val
+            if val >= self.count:
+                raise StopAsyncIteration
+            self.val += 1
+            return val
+
+    class async_zip(object):
+        def __init__(self, *largs):
+            self.nexts = [obj.__anext__ for obj in largs]
+
+        async def _accumulate(self, f, items, i):
+            items[i] = await f()
+
+        @aiter_compat
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            nexts = self.nexts
+            items = [
+                None,
+            ] * len(nexts)
+            got_stop = False
+
+            def handle(exc):
+                nonlocal got_stop
+                if isinstance(exc, StopAsyncIteration):
+                    got_stop = True
+                    return None
+                else:  # pragma: no cover
+                    return exc
+
+            with _core.MultiError.catch(handle):
+                async with _core.open_nursery() as nursery:
+                    for i, f in enumerate(nexts):
+                        nursery.start_soon(self._accumulate, f, items, i)
+
+            if got_stop:
+                raise StopAsyncIteration
+            return items
+
+    result = []
+    async for vals in async_zip(it(4), it(2)):
+        result.append(vals)
+    assert result == [[0, 0], [1, 1]]
+
+
 def test_contextvar_support():
     var = contextvars.ContextVar("test")
     var.set("before")