bpo-28893: Set __cause__ for errors in async iteration protocol (#407)

1st1 · 1st1 · commit dea5101ae101 · 2017-03-02T22:28:07.000-05:00
diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
@@ -1680,6 +1680,44 @@ async def foo():
                 warnings.simplefilter("error")
                 run_async(foo())
 
+    def test_for_11(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
+
+        with self.assertRaisesRegex(TypeError,
+                                    'an invalid object from __anext__') as c:
+            main().send(None)
+
+        err = c.exception
+        self.assertIsInstance(err.__cause__, ZeroDivisionError)
+
+    def test_for_12(self):
+        class F:
+            def __aiter__(self):
+                return self
+            def __await__(self):
+                1 / 0
+
+        async def main():
+            async for _ in F():
+                pass
+
+        with self.assertRaisesRegex(TypeError,
+                                    'an invalid object from __aiter__') as c:
+            main().send(None)
+
+        err = c.exception
+        self.assertIsInstance(err.__cause__, ZeroDivisionError)
+
     def test_for_tuple(self):
         class Done(Exception): pass
 
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -10,6 +10,9 @@ What's New in Python 3.6.1 release candidate 1?
 Core and Builtins
 -----------------
 
+- bpo-28893: Set correct __cause__ for errors about invalid awaitables
+  returned from __aiter__ and __anext__.
+
 - bpo-29683: Fixes to memory allocation in _PyCode_SetExtra.  Patch by
   Brian Coleman.
 
diff --git a/Python/ceval.c b/Python/ceval.c
@@ -1904,13 +1904,13 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
 
             awaitable = _PyCoro_GetAwaitableIter(iter);
             if (awaitable == NULL) {
-                SET_TOP(NULL);
-                PyErr_Format(
+                _PyErr_FormatFromCause(
                     PyExc_TypeError,
                     "'async for' received an invalid object "
                     "from __aiter__: %.100s",
                     Py_TYPE(iter)->tp_name);
 
+                SET_TOP(NULL);
                 Py_DECREF(iter);
                 goto error;
             } else {
@@ -1969,7 +1969,7 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
 
                 awaitable = _PyCoro_GetAwaitableIter(next_iter);
                 if (awaitable == NULL) {
-                    PyErr_Format(
+                    _PyErr_FormatFromCause(
                         PyExc_TypeError,
                         "'async for' received an invalid object "
                         "from __anext__: %.100s",
