Support return values in async yield from.

Author: ZeroIntensity

Include/cpython/pyerrors.h

@@ -73,6 +73,11 @@ typedef struct {
     PyObject *value;
 } PyStopIterationObject;
 
+typedef struct {
+    PyException_HEAD
+    PyObject *value;
+} PyStopAsyncIterationObject;
+
 typedef struct {
     PyException_HEAD
     PyObject *name;

Include/internal/pycore_opcode_metadata.h

@@ -720,8 +720,68 @@ SimpleExtendsException(PyExc_Exception, TypeError,
 /*
  *    StopAsyncIteration extends Exception
  */
-SimpleExtendsException(PyExc_Exception, StopAsyncIteration,
-                       "Signal the end from iterator.__anext__().");
+
+static PyMemberDef StopAsyncIteration_members[] = {
+    {"value", _Py_T_OBJECT, offsetof(PyStopAsyncIterationObject, value), 0,
+        PyDoc_STR("async generator return value")},
+    {NULL}  /* Sentinel */
+};
+
+static inline PyStopAsyncIterationObject *
+PyStopAsyncIterationObject_CAST(PyObject *self)
+{
+    assert(self != NULL);
+    assert(PyObject_TypeCheck(self, (PyTypeObject *)PyExc_StopAsyncIteration));
+    return (PyStopAsyncIterationObject *)self;
+}
+
+static int
+StopAsyncIteration_init(PyObject *op, PyObject *args, PyObject *kwds)
+{
+    Py_ssize_t size = PyTuple_GET_SIZE(args);
+
+    if (BaseException_init(op, args, kwds) < 0) {
+        return -1;
+    }
+    PyStopAsyncIterationObject *self = PyStopAsyncIterationObject_CAST(op);
+    Py_CLEAR(self->value);
+    PyObject *value;
+    if (size > 0) {
+        self->value = PyTuple_GET_ITEM(args, 0);
+    }
+    else {
+        self->value = Py_None;
+    };
+    return 0;
+}
+
+static int
+StopAsyncIteration_clear(PyObject *op)
+{
+    PyStopAsyncIterationObject *self = PyStopAsyncIterationObject_CAST(op);
+    Py_CLEAR(self->value);
+    return BaseException_clear(op);
+}
+
+static void
+StopAsyncIteration_dealloc(PyObject *self)
+{
+    _PyObject_GC_UNTRACK(self);
+    (void)StopAsyncIteration_clear(self);
+    Py_TYPE(self)->tp_free(self);
+}
+
+static int
+StopAsyncIteration_traverse(PyObject *op, visitproc visit, void *arg)
+{
+    PyStopAsyncIterationObject *self = PyStopAsyncIterationObject_CAST(op);
+    Py_VISIT(self->value);
+    return BaseException_traverse(op, visit, arg);
+}
+
+ComplexExtendsException(PyExc_Exception, StopAsyncIteration, StopAsyncIteration,
+                        0, 0, StopAsyncIteration_members, 0, 0, 0,
+                        "Signal the end from iterator.__anext__().");
 
 
 /*

Include/internal/pycore_uop_ids.h

@@ -307,8 +307,7 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult, int exc)
     /* If the generator just returned (as opposed to yielding), signal
      * that the generator is exhausted. */
     if (result) {
-        assert(result == Py_None || !PyAsyncGen_CheckExact(gen));
-        if (result == Py_None && !PyAsyncGen_CheckExact(gen) && !arg) {
+        if (result == Py_None && !arg) {
             /* Return NULL if called by gen_iternext() */
             Py_CLEAR(result);
         }
@@ -380,12 +379,14 @@ PyGen_am_send(PyObject *self, PyObject *arg, PyObject **result)
     return gen_send_ex(gen, arg, result);
 }
 
+int
+_PyAsyncGen_SetStopIterationValue(PyObject *value);
+
 static PyObject *
 gen_set_stop_iteration(PyGenObject *gen, PyObject *result)
 {
     if (PyAsyncGen_CheckExact(gen)) {
-        assert(result == Py_None);
-        PyErr_SetNone(PyExc_StopAsyncIteration);
+        _PyAsyncGen_SetStopIterationValue(result);
     }
     else if (result == Py_None) {
         PyErr_SetNone(PyExc_StopIteration);
@@ -776,24 +777,39 @@ gen_iternext(PyObject *self)
  * Returns 0 if StopIteration is set and -1 if any other exception is set.
  */
 int
-_PyGen_SetStopIterationValue(PyObject *value)
+_PyAnyGen_SetStopIterationValue(PyObject *exc_class, PyObject *value)
 {
+    assert(exc_class != NULL);
+    assert(PyType_Check(exc_class));
+    assert(value != NULL);
     assert(!PyErr_Occurred());
     // Construct an exception instance manually with PyObject_CallOneArg()
     // but use PyErr_SetRaisedException() instead of PyErr_SetObject() as
     // PyErr_SetObject(exc_type, value) has a fast path when 'value'
     // is a tuple, where the value of the StopIteration exception would be
     // set to 'value[0]' instead of 'value'.
     PyObject *exc = value == NULL
-        ? PyObject_CallNoArgs(PyExc_StopIteration)
-        : PyObject_CallOneArg(PyExc_StopIteration, value);
+        ? PyObject_CallNoArgs(exc_class)
+        : PyObject_CallOneArg(exc_class, value);
     if (exc == NULL) {
         return -1;
     }
     PyErr_SetRaisedException(exc /* stolen */);
     return 0;
 }
 
+int
+_PyGen_SetStopIterationValue(PyObject *value)
+{
+    return _PyAnyGen_SetStopIterationValue(PyExc_StopIteration, value);
+}
+
+int
+_PyAsyncGen_SetStopIterationValue(PyObject *value)
+{
+    return _PyAnyGen_SetStopIterationValue(PyExc_StopAsyncIteration, value);
+}
+
 /*
  *   If StopIteration exception is set, fetches its 'value'
  *   attribute if any, otherwise sets pvalue to None.
@@ -2645,6 +2661,7 @@ async_gen_yield_from_iternext(PyObject *op)
 {
     assert(op != NULL);
     _PyAsyncGenYieldFrom *self = _PyAsyncGenYieldFrom_CAST(op);
+    assert(self->agyf_iterator != NULL);
     PyObject *result = PyIter_Next(self->agyf_iterator);
     if (result == NULL) {
         return NULL;
@@ -2708,6 +2725,17 @@ _PyAsyncGenYieldFrom_New(PyThreadState *tstate, PyObject *iterator)
     if (yield_from == NULL) {
         return NULL;
     }
+    if (!PyIter_Check(iterator)) {
+        if (PyAsyncGen_CheckExact(iterator)) {
+            _PyErr_Format(tstate, PyExc_TypeError,
+                          "%T object is not iterable. Did you mean 'async yield from'?",
+                          iterator);
+        } else {
+            _PyErr_Format(tstate, PyExc_TypeError,
+                          "%T object is not iterable", iterator);
+        }
+        return NULL;
+    }
     yield_from->agyf_iterator = Py_NewRef(iterator);
     _PyObject_GC_TRACK((PyObject *)yield_from);
     return (PyObject *)yield_from;

Modules/_testinternalcapi/test_cases.c.h

@@ -2577,6 +2577,7 @@ static PyTypeObject* static_types[] = {
     &_PyAsyncGenASend_Type,
     &_PyAsyncGenAThrow_Type,
     &_PyAsyncGenWrappedValue_Type,
+    &_PyAsyncGenYieldFrom_Type,
     &_PyBufferWrapper_Type,
     &_PyContextTokenMissing_Type,
     &_PyCoroWrapper_Type,

Objects/exceptions.c

@@ -1583,14 +1583,17 @@ dummy_func(
             }
         }
 
-        inst(CLEANUP_ASYNC_THROW, (exc_value_st -- )) {
+        inst(CLEANUP_ASYNC_THROW, (iter, exc_value_st -- value)) {
             PyObject *exc_value = PyStackRef_AsPyObjectSteal(exc_value_st);
             assert(exc_value != NULL);
             assert(PyExceptionInstance_Check(exc_value));
             INPUTS_DEAD();
 
             int matches = PyErr_GivenExceptionMatches(exc_value, PyExc_StopAsyncIteration);
-            if (!matches) {
+            if (matches) {
+                value = PyStackRef_FromPyObjectNew(((PyStopAsyncIterationObject *)exc_value)->value);
+            }
+            else {
                 _PyErr_SetRaisedException(tstate, exc_value);
                 monitor_reraise(tstate, frame, this_instr);
                 goto exception_unwind;

Objects/genobject.c

@@ -2275,9 +2275,6 @@ codegen_return(compiler *c, stmt_ty s)
     if (!_PyST_IsFunctionLike(ste)) {
         return _PyCompile_Error(c, loc, "'return' outside function");
     }
-    if (s->v.Return.value != NULL && ste->ste_coroutine && ste->ste_generator) {
-        return _PyCompile_Error(c, loc, "'return' with value in async generator");
-    }
 
     if (preserve_tos) {
         VISIT(c, expr, s->v.Return.value);