patch

gen patch.patch

--- original_generator.c
+++ modified_generator.c
@@ -1,6 +1,7 @@
 /* Generator object implementation */

 #define _PY_INTERPRETER

+#include "pycore_exceptions.h" // For accessing exception attributes
 #include "Python.h"
 #include "pycore_call.h"          // _PyObject_CallNoArgs()
 #include "pycore_ceval.h"         // _PyEval_EvalFrame()
@@ -349,6 +350,31 @@
         /* No need to visit cr_origin, because it's just tuples/str/int, so can't
            participate in a reference cycle. */
         Py_VISIT(gen->gi_exc_state.exc_value);
         return 0;
+}
+
+/*
+ *   Set StopAsyncIteration with specified value. Value can be any object
+ *   or NULL.
+ *
+ *   Returns 0 if StopAsyncIteration is set and -1 if any other exception is set.
+ */
+int
+_PyAsyncGen_SetStopAsyncIterationValue(PyObject *value)
+{
+    PyObject *e;
+
+    if (value == NULL) {
+        PyErr_SetNone(PyExc_StopAsyncIteration);
+        return 0;
+    }
+
+    /* Construct StopAsyncIteration with the provided value */
+    e = PyObject_CallFunctionObjArgs(PyExc_StopAsyncIteration, value, NULL);
+    if (e == NULL) {
+        return -1;
+    }
+    PyErr_SetObject(PyExc_StopAsyncIteration, e);
+    Py_DECREF(e);
+    return 0;
+}
+
+/*
+ *   If StopAsyncIteration exception is set, fetches its 'value'
+ *   attribute if any, otherwise sets pvalue to None.
+ *
+ *   Returns 0 if StopAsyncIteration is set and value is fetched,
+ *   -1 otherwise.
+ */
+int
+_PyAsyncGen_FetchStopAsyncIterationValue(PyObject **pvalue)
+{
+    if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
+        PyObject *exc = PyErr_GetRaisedException();
+        PyObject *value = PyObject_GetAttr(exc, &_Py_ID(value));
+        Py_DECREF(exc);
+        if (value == NULL) {
+            if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
+                PyErr_Clear();
+                *pvalue = Py_NewRef(Py_None);
+                return 0;
+            }
+            return -1;
+        }
+        *pvalue = value;
+        return 0;
+    }
+    return -1;
 }

 PyCodeObject *
 PyGen_GetCode(PyGenObject *gen) {
@@ -1133,6 +1169,33 @@
     return result ? PYGEN_RETURN : PYGEN_ERROR;
 }

+/*
+ *   If StopAsyncIteration exception is set, fetches its 'value'
+ *   attribute if any, otherwise sets pvalue to None.
+ *
+ *   Returns 0 if StopAsyncIteration is set and value is fetched,
+ *   -1 otherwise.
+ */
+int
+_PyAsyncGen_FetchStopAsyncIterationValue(PyObject **pvalue)
+{
+    if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
+        PyObject *exc = PyErr_GetRaisedException();
+        PyObject *value = PyObject_GetAttr(exc, &_Py_ID(value));
+        Py_DECREF(exc);
+        if (value == NULL) {
+            if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
+                PyErr_Clear();
+                *pvalue = Py_NewRef(Py_None);
+                return 0;
+            }
+            return -1;
+        }
+        *pvalue = value;
+        return 0;
+    }
+    return -1;
+}
+
 static PySendResult
 gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
              int exc, int closing)
 {
     PyThreadState *tstate = _PyThreadState_GET();
     _PyInterpreterFrame *frame = &gen->gi_iframe;

     *presult = NULL;
     if (gen->gi_frame_state == FRAME_CREATED && arg && arg != Py_None) {
         const char *msg = "can't send non-None value to a "
                             "just-started generator";
         if (PyCoro_CheckExact(gen)) {
             msg = NON_INIT_CORO_MSG;
         }
         else if (PyAsyncGen_CheckExact(gen)) {
             msg = "can't send non-None value to a "
                     "just-started async generator";
         }
         PyErr_SetString(PyExc_TypeError, msg);
         return PYGEN_ERROR;
     }
     if (gen->gi_frame_state == FRAME_EXECUTING) {
         const char *msg = "generator already executing";
         if (PyCoro_CheckExact(gen)) {
             msg = "coroutine already executing";
         }
         else if (PyAsyncGen_CheckExact(gen)) {
             msg = "async generator already executing";
         }
         PyErr_SetString(PyExc_ValueError, msg);
         return PYGEN_ERROR;
     }
     if (FRAME_STATE_FINISHED(gen->gi_frame_state)) {
         if (PyCoro_CheckExact(gen) && !closing) {
             /* `gen` is an exhausted coroutine: raise an error,
                except when called from gen_close(), which should
                always be a silent method. */
             PyErr_SetString(
                 PyExc_RuntimeError,
                 "cannot reuse already awaited coroutine");
         }
         else if (arg && !exc) {
             /* `gen` is an exhausted generator:
                only return value if called from send(). */
             *presult = Py_NewRef(Py_None);
             return PYGEN_RETURN;
         }
         return PYGEN_ERROR;
     }

     assert((gen->gi_frame_state == FRAME_CREATED) ||
            FRAME_STATE_SUSPENDED(gen->gi_frame_state));

     /* Push arg onto the frame's value stack */
     PyObject *arg_obj = arg ? arg : Py_None;
     _PyFrame_StackPush(frame, PyStackRef_FromPyObjectNew(arg_obj));

     _PyErr_StackItem *prev_exc_info = tstate->exc_info;
     gen->gi_exc_state.previous_item = prev_exc_info;
     tstate->exc_info = &gen->gi_exc_state;

     if (exc) {
         assert(_PyErr_Occurred(tstate));
         _PyErr_ChainStackItem();
     }

     gen->gi_frame_state = FRAME_EXECUTING;
     EVAL_CALL_STAT_INC(EVAL_CALL_GENERATOR);
     PyObject *result = _PyEval_EvalFrame(tstate, frame, exc);
     assert(tstate->exc_info == prev_exc_info);
     assert(gen->gi_exc_state.previous_item == NULL);
     assert(gen->gi_frame_state != FRAME_EXECUTING);
     assert(frame->previous == NULL);

     /* If the generator just returned (as opposed to yielding), signal
      * that the generator is exhausted. */
     if (result) {
         if (FRAME_STATE_SUSPENDED(gen->gi_frame_state)) {
             *presult = result;
             return PYGEN_NEXT;
         }
         if (PyAsyncGen_CheckExact(gen)) {
             /* For async generators, set StopAsyncIteration with the return value */
             if (_PyAsyncGen_SetStopAsyncIterationValue(result) < 0) {
                 /* Failed to set StopAsyncIteration */
                 return PYGEN_ERROR;
             }
             Py_DECREF(result);
             *presult = NULL;
             return PYGEN_RETURN;
         }
         assert(result == Py_None || !PyAsyncGen_CheckExact(gen));
         if (result == Py_None && !PyAsyncGen_CheckExact(gen) && !arg) {
             /* Return NULL if called by gen_iternext() */
             Py_CLEAR(result);
         }
     }
     else {
         assert(!PyErr_ExceptionMatches(PyExc_StopIteration));
         assert(!PyAsyncGen_CheckExact(gen) ||
             !PyErr_ExceptionMatches(PyExc_StopAsyncIteration));
     }

     assert(gen->gi_exc_state.exc_value == NULL);
     assert(gen->gi_frame_state == FRAME_CLEARED);
     *presult = result;
     return result ? PYGEN_RETURN : PYGEN_ERROR;
 }

@@ -1354,6 +1491,31 @@
     return result ? PYGEN_RETURN : PYGEN_ERROR;
 }

 /* Coroutine Object */

+/*
+ *   If StopAsyncIteration exception is set, fetches its 'value'
+ *   attribute if any, otherwise sets pvalue to None.
+ *
+ *   Returns 0 if StopAsyncIteration is set and value is fetched,
+ *   -1 otherwise.
+ */
+int
+_PyAsyncGen_FetchStopAsyncIterationValue(PyObject **pvalue)
+{
+    if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
+        PyObject *exc = PyErr_GetRaisedException();
+        PyObject *value = PyObject_GetAttr(exc, &_Py_ID(value));
+        Py_DECREF(exc);
+        if (value == NULL) {
+            if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
+                PyErr_Clear();
+                *pvalue = Py_NewRef(Py_None);
+                return 0;
+            }
+            return -1;
+        }
+        *pvalue = value;
+        return 0;
+    }
+    return -1;
+}
+
 static PyObject *
 gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
 {
     PyObject *result;
     if (gen_send_ex2(gen, arg, &result, exc, closing) == PYGEN_RETURN) {
         if (PyAsyncGen_CheckExact(gen)) {
             /* For async generators, retrieve the return value from StopAsyncIteration */
-            PyObject *exc = PyErr_GetRaisedException();
-            if (_PyGen_FetchStopIterationValue(&result) == 0) {
-                /* Attach the value to StopAsyncIteration */
-                PyErr_SetObject(PyExc_StopAsyncIteration, result);
-                Py_DECREF(result);
+            /* Fetch the value from StopAsyncIteration */
+            if (_PyAsyncGen_FetchStopAsyncIterationValue(&result) == 0) {
+                /* The value is already set in the exception */
+                /* Return the value to the caller */
+                return result;
+            }
+            /* If fetching the value failed, propagate the exception */
+            return NULL;
         }
         else if (result == Py_None) {
             PyErr_SetNone(PyExc_StopIteration);
             Py_CLEAR(result);
         }
         else {
             _PyGen_SetStopIterationValue(result);
         }
         Py_CLEAR(result);
     }
     return result;
 }
@@ -1663,6 +1767,23 @@
     return NULL;
 }

+/*
+ *   Finalize function for async generators to handle return values.
+ */
+void
+_PyAsyncGen_Finalize(PyObject *self)
+{
+    PyAsyncGenObject *ag = (PyAsyncGenObject *)self;
+
+    if (FRAME_STATE_FINISHED(ag->ag_frame_state)) {
+        /* Async generator isn't paused, so no need to close */
+        return;
+    }
+
+    /* Save the current exception, if any. */
+    PyObject *exc = PyErr_GetRaisedException();
+
+    /* Close the generator */
+    PyObject *res = gen_close((PyObject*)ag, NULL);
+    if (res == NULL) {
+        if (PyErr_Occurred()) {
+            PyErr_WriteUnraisable(self);
+        }
+    }
+    else {
+        Py_DECREF(res);
+    }
+
+    /* Restore the saved exception. */
+    PyErr_SetRaisedException(exc);
+}
+
 void
 _PyGen_Finalize(PyObject *self)
 {
@@ -1748,6 +1890,7 @@
         }
     }

+    if (PyAsyncGen_CheckExact(self)) {
         /* Save the current exception, if any. */
         PyObject *exc = PyErr_GetRaisedException();

@@ -1758,6 +1901,16 @@
         PyErr_SetRaisedException(exc);
     }
     else {
+        if (PyAsyncGen_CheckExact(self)) {
+            /* Call the specific async generator finalizer */
+            _PyAsyncGen_Finalize(self);
+        }
+
+        /* Existing finalization logic for regular generators and coroutines */
         PyObject *res = gen_close((PyObject*)gen, NULL);
         if (res == NULL) {
             if (PyErr_Occurred()) {
                 PyErr_WriteUnraisable(self);
             }
             else {
                 Py_DECREF(res);
             }
         }
     }
 }
@@ -2333,6 +2456,32 @@
         /* Normal finalization steps */
         Py_CLEAR(ags->ags_gen);
         Py_CLEAR(ags->ags_sendval);
+
+        /* Finalize return value if stored (optional) */
+        if (ags->ags_gen->ag_return_value) {
+            Py_CLEAR(ags->ags_gen->ag_return_value);
+        }
+
         _PyObject_GC_UNTRACK((PyObject*)ags);
         PyObject_GC_Del(ags);
     }

+/* Async Generator AThrow awaitable modifications */
+static void
+async_gen_athrow_finalize(PyObject *self)
+{
+    PyAsyncGenAThrow *agt = _PyAsyncGenAThrow_CAST(self);
+    if (agt->agt_state == AWAITABLE_STATE_INIT) {
+        PyObject *method = agt->agt_args ? &_Py_ID(athrow) : &_Py_ID(aclose);
+        _PyErr_WarnUnawaitedAgenMethod(agt->agt_gen, method);
+    }
+}
+
+/*
+ *   Modify async_gen_asend_send to handle return values.
+ */
+static PyObject *
+async_gen_asend_send(PyObject *self, PyObject *arg)
+{
+    PyAsyncGenASend *o = _PyAsyncGenASend_CAST(self);
+    if (o->ags_state == AWAITABLE_STATE_CLOSED) {
+        PyErr_SetString(
+            PyExc_RuntimeError,
+            "cannot reuse already awaited __anext__()/asend()");
+        return NULL;
+    }
+
+    if (o->ags_state == AWAITABLE_STATE_INIT) {
+        if (o->ags_gen->ag_running_async) {
+            o->ags_state = AWAITABLE_STATE_CLOSED;
+            PyErr_SetString(
+                PyExc_RuntimeError,
+                "anext(): asynchronous generator is already running");
+            return NULL;
+        }
+
+        if (arg == NULL || arg == Py_None) {
+            arg = o->ags_sendval;
+        }
+        o->ags_state = AWAITABLE_STATE_ITER;
+    }
+
+    o->ags_gen->ag_running_async = 1;
+    PyObject *result = gen_send((PyObject*)o->ags_gen, arg);
+    result = async_gen_unwrap_value(o->ags_gen, result);
+
+    if (result == NULL) {
+        o->ags_state = AWAITABLE_STATE_CLOSED;
+    }
+
+    return result;
+}
+
 /*
  *   Set StopIteration with specified value.  Value can be arbitrary object
  *   or NULL.
@@ -2450,6 +2623,13 @@
         return NULL;
     }
 
+/* Fetch the value from StopAsyncIteration */
+static int
+async_gen_unwrap_stopasynciteration(PyObject *result)
+{
+    if (_PyAsyncGen_FetchStopAsyncIterationValue(&result) == 0) {
+        /* The value is already set in the exception */
+        return 0;
+    }
+    return -1;
+}
+
 static PyObject *
 async_gen_aclose(PyAsyncGenObject *o, PyObject *arg)
 {
@@ -2548,6 +2738,7 @@
         o->agt_gen->ag_running_async = 0;
         o->agt_state = AWAITABLE_STATE_CLOSED;
     }

+/* Async generator throw modifications */
 static PyObject *
 async_gen_athrow(PyAsyncGenObject *o, PyObject *args)
 {