Issue #27243: Fix __aiter__ protocol

1 files changed, 38 insertions, 2 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index 3758b09..3d69038 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1933,8 +1933,9 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             PyObject *obj = TOP();
             PyTypeObject *type = Py_TYPE(obj);
 
-            if (type->tp_as_async != NULL)
+            if (type->tp_as_async != NULL) {
                 getter = type->tp_as_async->am_aiter;
+            }
 
             if (getter != NULL) {
                 iter = (*getter)(obj);
@@ -1955,6 +1956,27 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
                 goto error;
             }
 
+            if (Py_TYPE(iter)->tp_as_async != NULL &&
+                    Py_TYPE(iter)->tp_as_async->am_anext != NULL) {
+
+                /* Starting with CPython 3.5.2 __aiter__ should return
+                   asynchronous iterators directly (not awaitables that
+                   resolve to asynchronous iterators.)
+
+                   Therefore, we check if the object that was returned
+                   from __aiter__ has an __anext__ method.  If it does,
+                   we wrap it in an awaitable that resolves to `iter`.
+
+                   See http://bugs.python.org/issue27243 for more
+                   details.
+                */
+
+                PyObject *wrapper = _PyAIterWrapper_New(iter);
+                Py_DECREF(iter);
+                SET_TOP(wrapper);
+                DISPATCH();
+            }
+
             awaitable = _PyCoro_GetAwaitableIter(iter);
             if (awaitable == NULL) {
                 SET_TOP(NULL);
@@ -1966,9 +1988,23 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 
                 Py_DECREF(iter);
                 goto error;
-            } else
+            } else {
                 Py_DECREF(iter);
 
+                if (PyErr_WarnFormat(
+                        PyExc_PendingDeprecationWarning, 1,
+                        "'%.100s' implements legacy __aiter__ protocol; "
+                        "__aiter__ should return an asynchronous "
+                        "iterator, not awaitable",
+                        type->tp_name))
+                {
+                    /* Warning was converted to an error. */
+                    Py_DECREF(awaitable);
+                    SET_TOP(NULL);
+                    goto error;
+                }
+            }
+
             SET_TOP(awaitable);
             DISPATCH();
         }