Skip to content

Commit b6237c3

Browse files
gh-117852: fix argument checking of async_generator.athrow (#134868)
Co-authored-by: Kumar Aditya <kumaraditya@python.org>
1 parent 9c72658 commit b6237c3

File tree

3 files changed

+38
-21
lines changed

3 files changed

+38
-21
lines changed

Lib/test/test_asyncgen.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2021,6 +2021,15 @@ async def gen():
20212021
g.athrow(RuntimeError)
20222022
gc_collect()
20232023

2024+
def test_athrow_throws_immediately(self):
2025+
async def gen():
2026+
yield 1
2027+
2028+
g = gen()
2029+
msg = "athrow expected at least 1 argument, got 0"
2030+
with self.assertRaisesRegex(TypeError, msg):
2031+
g.athrow()
2032+
20242033
def test_aclose(self):
20252034
async def gen():
20262035
yield 1
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Fix argument checking of :meth:`~agen.athrow`.

Objects/genobject.c

Lines changed: 28 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1451,7 +1451,9 @@ typedef struct PyAsyncGenAThrow {
14511451

14521452
/* Can be NULL, when in the "aclose()" mode
14531453
(equivalent of "athrow(GeneratorExit)") */
1454-
PyObject *agt_args;
1454+
PyObject *agt_typ;
1455+
PyObject *agt_tb;
1456+
PyObject *agt_val;
14551457

14561458
AwaitableState agt_state;
14571459
} PyAsyncGenAThrow;
@@ -2078,7 +2080,9 @@ async_gen_athrow_dealloc(PyObject *self)
20782080

20792081
_PyObject_GC_UNTRACK(self);
20802082
Py_CLEAR(agt->agt_gen);
2081-
Py_CLEAR(agt->agt_args);
2083+
Py_XDECREF(agt->agt_typ);
2084+
Py_XDECREF(agt->agt_tb);
2085+
Py_XDECREF(agt->agt_val);
20822086
PyObject_GC_Del(self);
20832087
}
20842088

@@ -2088,7 +2092,9 @@ async_gen_athrow_traverse(PyObject *self, visitproc visit, void *arg)
20882092
{
20892093
PyAsyncGenAThrow *agt = _PyAsyncGenAThrow_CAST(self);
20902094
Py_VISIT(agt->agt_gen);
2091-
Py_VISIT(agt->agt_args);
2095+
Py_VISIT(agt->agt_typ);
2096+
Py_VISIT(agt->agt_tb);
2097+
Py_VISIT(agt->agt_val);
20922098
return 0;
20932099
}
20942100

@@ -2116,7 +2122,7 @@ async_gen_athrow_send(PyObject *self, PyObject *arg)
21162122
if (o->agt_state == AWAITABLE_STATE_INIT) {
21172123
if (o->agt_gen->ag_running_async) {
21182124
o->agt_state = AWAITABLE_STATE_CLOSED;
2119-
if (o->agt_args == NULL) {
2125+
if (o->agt_typ == NULL) {
21202126
PyErr_SetString(
21212127
PyExc_RuntimeError,
21222128
"aclose(): asynchronous generator is already running");
@@ -2143,7 +2149,7 @@ async_gen_athrow_send(PyObject *self, PyObject *arg)
21432149
o->agt_state = AWAITABLE_STATE_ITER;
21442150
o->agt_gen->ag_running_async = 1;
21452151

2146-
if (o->agt_args == NULL) {
2152+
if (o->agt_typ == NULL) {
21472153
/* aclose() mode */
21482154
o->agt_gen->ag_closed = 1;
21492155

@@ -2157,19 +2163,10 @@ async_gen_athrow_send(PyObject *self, PyObject *arg)
21572163
goto yield_close;
21582164
}
21592165
} else {
2160-
PyObject *typ;
2161-
PyObject *tb = NULL;
2162-
PyObject *val = NULL;
2163-
2164-
if (!PyArg_UnpackTuple(o->agt_args, "athrow", 1, 3,
2165-
&typ, &val, &tb)) {
2166-
return NULL;
2167-
}
2168-
21692166
retval = _gen_throw((PyGenObject *)gen,
21702167
0, /* Do not close generator when
21712168
PyExc_GeneratorExit is passed */
2172-
typ, val, tb);
2169+
o->agt_typ, o->agt_val, o->agt_tb);
21732170
retval = async_gen_unwrap_value(o->agt_gen, retval);
21742171
}
21752172
if (retval == NULL) {
@@ -2181,7 +2178,7 @@ async_gen_athrow_send(PyObject *self, PyObject *arg)
21812178
assert(o->agt_state == AWAITABLE_STATE_ITER);
21822179

21832180
retval = gen_send((PyObject *)gen, arg);
2184-
if (o->agt_args) {
2181+
if (o->agt_typ) {
21852182
return async_gen_unwrap_value(o->agt_gen, retval);
21862183
} else {
21872184
/* aclose() mode */
@@ -2212,7 +2209,7 @@ async_gen_athrow_send(PyObject *self, PyObject *arg)
22122209
if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) ||
22132210
PyErr_ExceptionMatches(PyExc_GeneratorExit))
22142211
{
2215-
if (o->agt_args == NULL) {
2212+
if (o->agt_typ == NULL) {
22162213
/* when aclose() is called we don't want to propagate
22172214
StopAsyncIteration or GeneratorExit; just raise
22182215
StopIteration, signalling that this 'aclose()' await
@@ -2241,7 +2238,7 @@ async_gen_athrow_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
22412238
if (o->agt_state == AWAITABLE_STATE_INIT) {
22422239
if (o->agt_gen->ag_running_async) {
22432240
o->agt_state = AWAITABLE_STATE_CLOSED;
2244-
if (o->agt_args == NULL) {
2241+
if (o->agt_typ == NULL) {
22452242
PyErr_SetString(
22462243
PyExc_RuntimeError,
22472244
"aclose(): asynchronous generator is already running");
@@ -2259,7 +2256,7 @@ async_gen_athrow_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
22592256
}
22602257

22612258
PyObject *retval = gen_throw((PyObject*)o->agt_gen, args, nargs);
2262-
if (o->agt_args) {
2259+
if (o->agt_typ) {
22632260
retval = async_gen_unwrap_value(o->agt_gen, retval);
22642261
if (retval == NULL) {
22652262
o->agt_gen->ag_running_async = 0;
@@ -2334,7 +2331,7 @@ async_gen_athrow_finalize(PyObject *op)
23342331
{
23352332
PyAsyncGenAThrow *o = (PyAsyncGenAThrow*)op;
23362333
if (o->agt_state == AWAITABLE_STATE_INIT) {
2337-
PyObject *method = o->agt_args ? &_Py_ID(athrow) : &_Py_ID(aclose);
2334+
PyObject *method = o->agt_typ ? &_Py_ID(athrow) : &_Py_ID(aclose);
23382335
_PyErr_WarnUnawaitedAgenMethod(o->agt_gen, method);
23392336
}
23402337
}
@@ -2403,13 +2400,23 @@ PyTypeObject _PyAsyncGenAThrow_Type = {
24032400
static PyObject *
24042401
async_gen_athrow_new(PyAsyncGenObject *gen, PyObject *args)
24052402
{
2403+
PyObject *typ = NULL;
2404+
PyObject *tb = NULL;
2405+
PyObject *val = NULL;
2406+
if (args && !PyArg_UnpackTuple(args, "athrow", 1, 3, &typ, &val, &tb)) {
2407+
return NULL;
2408+
}
2409+
24062410
PyAsyncGenAThrow *o;
24072411
o = PyObject_GC_New(PyAsyncGenAThrow, &_PyAsyncGenAThrow_Type);
24082412
if (o == NULL) {
24092413
return NULL;
24102414
}
24112415
o->agt_gen = (PyAsyncGenObject*)Py_NewRef(gen);
2412-
o->agt_args = Py_XNewRef(args);
2416+
o->agt_typ = Py_XNewRef(typ);
2417+
o->agt_tb = Py_XNewRef(tb);
2418+
o->agt_val = Py_XNewRef(val);
2419+
24132420
o->agt_state = AWAITABLE_STATE_INIT;
24142421
_PyObject_GC_TRACK((PyObject*)o);
24152422
return (PyObject*)o;

0 commit comments

Comments
 (0)