[pypy-commit] pypy py3.6: Remember the __cause__ when an exception is raised in __await__.

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3.6
Changeset: r93958:ab182b33a9c1
Date: 2018-03-05 09:26 +0100
http://bitbucket.org/pypy/pypy/changeset/ab182b33a9c1/

Log:	Remember the __cause__ when an exception is raised in __await__.

diff --git a/pypy/interpreter/pyopcode.py b/pypy/interpreter/pyopcode.py
--- a/pypy/interpreter/pyopcode.py
+++ b/pypy/interpreter/pyopcode.py
@@ -1596,14 +1596,15 @@
             try:
                 w_awaitable = get_awaitable_iter(space, w_iter)
             except OperationError as e:
-                # yay! get_awaitable_iter() carefully builds a useful
-                # error message, but here we're eating *all errors*
-                # to replace it with a generic one.
                 if e.async(space):
                     raise
-                raise oefmt(space.w_TypeError,
+                new_error = oefmt(space.w_TypeError,
                             "'async for' received an invalid object "
                             "from __aiter__: %T", w_iter)
+                e.normalize_exception(space)
+                new_error.normalize_exception(space)
+                new_error.set_cause(space, e.get_w_value(space))
+                raise new_error
             space.warn(space.newtext(
                 "'%s' implements legacy __aiter__ protocol; "
                 "__aiter__ should return an asynchronous "
@@ -1628,14 +1629,15 @@
         try:
             w_awaitable = get_awaitable_iter(space, w_next_iter)
         except OperationError as e:
-            # yay! get_awaitable_iter() carefully builds a useful
-            # error message, but here we're eating *all errors*
-            # to replace it with a generic one.
             if e.async(space):
                 raise
-            raise oefmt(space.w_TypeError,
+            new_error = oefmt(space.w_TypeError,
                         "'async for' received an invalid object "
                         "from __anext__: %T", w_next_iter)
+            e.normalize_exception(space)
+            new_error.normalize_exception(space)
+            new_error.set_cause(space, e.get_w_value(space))
+            raise new_error
         self.pushvalue(w_awaitable)
 
     def FORMAT_VALUE(self, oparg, next_instr):
diff --git a/pypy/interpreter/test/test_coroutine.py b/pypy/interpreter/test/test_coroutine.py
--- a/pypy/interpreter/test/test_coroutine.py
+++ b/pypy/interpreter/test/test_coroutine.py
@@ -75,6 +75,24 @@
         assert next(cr.__await__()) == 20
         """
 
+    def test_for_error_cause(self): """
+        class F:
+            def __aiter__(self):
+                return self
+            def __anext__(self):
+                return self
+            def __await__(self):
+                1 / 0
+
+        async def main():
+            async for _ in F():
+                pass
+
+        c = raises(TypeError, main().send, None)
+        assert 'an invalid object from __anext__' in c.value.args[0], c.value
+        assert isinstance(c.value.__cause__, ZeroDivisionError)
+        """
+
     def test_set_coroutine_wrapper(self): """
         import sys
         async def f():