• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Generator object implementation */
2 
3 #include "Python.h"
4 #include "pycore_ceval.h"         // _PyEval_EvalFrame()
5 #include "pycore_object.h"
6 #include "pycore_pyerrors.h"      // _PyErr_ClearExcState()
7 #include "pycore_pystate.h"       // _PyThreadState_GET()
8 #include "frameobject.h"
9 #include "structmember.h"         // PyMemberDef
10 #include "opcode.h"
11 
12 static PyObject *gen_close(PyGenObject *, PyObject *);
13 static PyObject *async_gen_asend_new(PyAsyncGenObject *, PyObject *);
14 static PyObject *async_gen_athrow_new(PyAsyncGenObject *, PyObject *);
15 
16 static const char *NON_INIT_CORO_MSG = "can't send non-None value to a "
17                                  "just-started coroutine";
18 
19 static const char *ASYNC_GEN_IGNORED_EXIT_MSG =
20                                  "async generator ignored GeneratorExit";
21 
22 static inline int
exc_state_traverse(_PyErr_StackItem * exc_state,visitproc visit,void * arg)23 exc_state_traverse(_PyErr_StackItem *exc_state, visitproc visit, void *arg)
24 {
25     Py_VISIT(exc_state->exc_type);
26     Py_VISIT(exc_state->exc_value);
27     Py_VISIT(exc_state->exc_traceback);
28     return 0;
29 }
30 
31 static int
gen_traverse(PyGenObject * gen,visitproc visit,void * arg)32 gen_traverse(PyGenObject *gen, visitproc visit, void *arg)
33 {
34     Py_VISIT((PyObject *)gen->gi_frame);
35     Py_VISIT(gen->gi_code);
36     Py_VISIT(gen->gi_name);
37     Py_VISIT(gen->gi_qualname);
38     /* No need to visit cr_origin, because it's just tuples/str/int, so can't
39        participate in a reference cycle. */
40     return exc_state_traverse(&gen->gi_exc_state, visit, arg);
41 }
42 
43 void
_PyGen_Finalize(PyObject * self)44 _PyGen_Finalize(PyObject *self)
45 {
46     PyGenObject *gen = (PyGenObject *)self;
47     PyObject *res = NULL;
48     PyObject *error_type, *error_value, *error_traceback;
49 
50     if (gen->gi_frame == NULL ||  _PyFrameHasCompleted(gen->gi_frame)) {
51         /* Generator isn't paused, so no need to close */
52         return;
53     }
54 
55     if (PyAsyncGen_CheckExact(self)) {
56         PyAsyncGenObject *agen = (PyAsyncGenObject*)self;
57         PyObject *finalizer = agen->ag_finalizer;
58         if (finalizer && !agen->ag_closed) {
59             /* Save the current exception, if any. */
60             PyErr_Fetch(&error_type, &error_value, &error_traceback);
61 
62             res = PyObject_CallOneArg(finalizer, self);
63 
64             if (res == NULL) {
65                 PyErr_WriteUnraisable(self);
66             } else {
67                 Py_DECREF(res);
68             }
69             /* Restore the saved exception. */
70             PyErr_Restore(error_type, error_value, error_traceback);
71             return;
72         }
73     }
74 
75     /* Save the current exception, if any. */
76     PyErr_Fetch(&error_type, &error_value, &error_traceback);
77 
78     /* If `gen` is a coroutine, and if it was never awaited on,
79        issue a RuntimeWarning. */
80     if (gen->gi_code != NULL &&
81         ((PyCodeObject *)gen->gi_code)->co_flags & CO_COROUTINE &&
82         gen->gi_frame->f_lasti == -1)
83     {
84         _PyErr_WarnUnawaitedCoroutine((PyObject *)gen);
85     }
86     else {
87         res = gen_close(gen, NULL);
88     }
89 
90     if (res == NULL) {
91         if (PyErr_Occurred()) {
92             PyErr_WriteUnraisable(self);
93         }
94     }
95     else {
96         Py_DECREF(res);
97     }
98 
99     /* Restore the saved exception. */
100     PyErr_Restore(error_type, error_value, error_traceback);
101 }
102 
103 static void
gen_dealloc(PyGenObject * gen)104 gen_dealloc(PyGenObject *gen)
105 {
106     PyObject *self = (PyObject *) gen;
107 
108     _PyObject_GC_UNTRACK(gen);
109 
110     if (gen->gi_weakreflist != NULL)
111         PyObject_ClearWeakRefs(self);
112 
113     _PyObject_GC_TRACK(self);
114 
115     if (PyObject_CallFinalizerFromDealloc(self))
116         return;                     /* resurrected.  :( */
117 
118     _PyObject_GC_UNTRACK(self);
119     if (PyAsyncGen_CheckExact(gen)) {
120         /* We have to handle this case for asynchronous generators
121            right here, because this code has to be between UNTRACK
122            and GC_Del. */
123         Py_CLEAR(((PyAsyncGenObject*)gen)->ag_finalizer);
124     }
125     if (gen->gi_frame != NULL) {
126         gen->gi_frame->f_gen = NULL;
127         Py_CLEAR(gen->gi_frame);
128     }
129     if (((PyCodeObject *)gen->gi_code)->co_flags & CO_COROUTINE) {
130         Py_CLEAR(((PyCoroObject *)gen)->cr_origin);
131     }
132     Py_CLEAR(gen->gi_code);
133     Py_CLEAR(gen->gi_name);
134     Py_CLEAR(gen->gi_qualname);
135     _PyErr_ClearExcState(&gen->gi_exc_state);
136     PyObject_GC_Del(gen);
137 }
138 
139 static PySendResult
gen_send_ex2(PyGenObject * gen,PyObject * arg,PyObject ** presult,int exc,int closing)140 gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
141              int exc, int closing)
142 {
143     PyThreadState *tstate = _PyThreadState_GET();
144     PyFrameObject *f = gen->gi_frame;
145     PyObject *result;
146 
147     *presult = NULL;
148     if (f != NULL && f->f_lasti < 0 && arg && arg != Py_None) {
149         const char *msg = "can't send non-None value to a "
150                             "just-started generator";
151         if (PyCoro_CheckExact(gen)) {
152             msg = NON_INIT_CORO_MSG;
153         }
154         else if (PyAsyncGen_CheckExact(gen)) {
155             msg = "can't send non-None value to a "
156                     "just-started async generator";
157         }
158         PyErr_SetString(PyExc_TypeError, msg);
159         return PYGEN_ERROR;
160     }
161     if (f != NULL && _PyFrame_IsExecuting(f)) {
162         const char *msg = "generator already executing";
163         if (PyCoro_CheckExact(gen)) {
164             msg = "coroutine already executing";
165         }
166         else if (PyAsyncGen_CheckExact(gen)) {
167             msg = "async generator already executing";
168         }
169         PyErr_SetString(PyExc_ValueError, msg);
170         return PYGEN_ERROR;
171     }
172     if (f == NULL || _PyFrameHasCompleted(f)) {
173         if (PyCoro_CheckExact(gen) && !closing) {
174             /* `gen` is an exhausted coroutine: raise an error,
175                except when called from gen_close(), which should
176                always be a silent method. */
177             PyErr_SetString(
178                 PyExc_RuntimeError,
179                 "cannot reuse already awaited coroutine");
180         }
181         else if (arg && !exc) {
182             /* `gen` is an exhausted generator:
183                only return value if called from send(). */
184             *presult = Py_None;
185             Py_INCREF(*presult);
186             return PYGEN_RETURN;
187         }
188         return PYGEN_ERROR;
189     }
190 
191     assert(_PyFrame_IsRunnable(f));
192     assert(f->f_lasti >= 0 || ((unsigned char *)PyBytes_AS_STRING(f->f_code->co_code))[0] == GEN_START);
193     /* Push arg onto the frame's value stack */
194     result = arg ? arg : Py_None;
195     Py_INCREF(result);
196     gen->gi_frame->f_valuestack[gen->gi_frame->f_stackdepth] = result;
197     gen->gi_frame->f_stackdepth++;
198 
199     /* Generators always return to their most recent caller, not
200      * necessarily their creator. */
201     Py_XINCREF(tstate->frame);
202     assert(f->f_back == NULL);
203     f->f_back = tstate->frame;
204 
205     gen->gi_exc_state.previous_item = tstate->exc_info;
206     tstate->exc_info = &gen->gi_exc_state;
207 
208     if (exc) {
209         assert(_PyErr_Occurred(tstate));
210         _PyErr_ChainStackItem(NULL);
211     }
212 
213     result = _PyEval_EvalFrame(tstate, f, exc);
214     tstate->exc_info = gen->gi_exc_state.previous_item;
215     gen->gi_exc_state.previous_item = NULL;
216 
217     /* Don't keep the reference to f_back any longer than necessary.  It
218      * may keep a chain of frames alive or it could create a reference
219      * cycle. */
220     assert(f->f_back == tstate->frame);
221     Py_CLEAR(f->f_back);
222 
223     /* If the generator just returned (as opposed to yielding), signal
224      * that the generator is exhausted. */
225     if (result) {
226         if (!_PyFrameHasCompleted(f)) {
227             *presult = result;
228             return PYGEN_NEXT;
229         }
230         assert(result == Py_None || !PyAsyncGen_CheckExact(gen));
231         if (result == Py_None && !PyAsyncGen_CheckExact(gen) && !arg) {
232             /* Return NULL if called by gen_iternext() */
233             Py_CLEAR(result);
234         }
235     }
236     else {
237         if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
238             const char *msg = "generator raised StopIteration";
239             if (PyCoro_CheckExact(gen)) {
240                 msg = "coroutine raised StopIteration";
241             }
242             else if (PyAsyncGen_CheckExact(gen)) {
243                 msg = "async generator raised StopIteration";
244             }
245             _PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
246         }
247         else if (PyAsyncGen_CheckExact(gen) &&
248                 PyErr_ExceptionMatches(PyExc_StopAsyncIteration))
249         {
250             /* code in `gen` raised a StopAsyncIteration error:
251                raise a RuntimeError.
252             */
253             const char *msg = "async generator raised StopAsyncIteration";
254             _PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
255         }
256     }
257 
258     /* generator can't be rerun, so release the frame */
259     /* first clean reference cycle through stored exception traceback */
260     _PyErr_ClearExcState(&gen->gi_exc_state);
261     gen->gi_frame->f_gen = NULL;
262     gen->gi_frame = NULL;
263     Py_DECREF(f);
264 
265     *presult = result;
266     return result ? PYGEN_RETURN : PYGEN_ERROR;
267 }
268 
269 static PySendResult
PyGen_am_send(PyGenObject * gen,PyObject * arg,PyObject ** result)270 PyGen_am_send(PyGenObject *gen, PyObject *arg, PyObject **result)
271 {
272     return gen_send_ex2(gen, arg, result, 0, 0);
273 }
274 
275 static PyObject *
gen_send_ex(PyGenObject * gen,PyObject * arg,int exc,int closing)276 gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
277 {
278     PyObject *result;
279     if (gen_send_ex2(gen, arg, &result, exc, closing) == PYGEN_RETURN) {
280         if (PyAsyncGen_CheckExact(gen)) {
281             assert(result == Py_None);
282             PyErr_SetNone(PyExc_StopAsyncIteration);
283         }
284         else if (result == Py_None) {
285             PyErr_SetNone(PyExc_StopIteration);
286         }
287         else {
288             _PyGen_SetStopIterationValue(result);
289         }
290         Py_CLEAR(result);
291     }
292     return result;
293 }
294 
295 PyDoc_STRVAR(send_doc,
296 "send(arg) -> send 'arg' into generator,\n\
297 return next yielded value or raise StopIteration.");
298 
299 static PyObject *
gen_send(PyGenObject * gen,PyObject * arg)300 gen_send(PyGenObject *gen, PyObject *arg)
301 {
302     return gen_send_ex(gen, arg, 0, 0);
303 }
304 
305 PyDoc_STRVAR(close_doc,
306 "close() -> raise GeneratorExit inside generator.");
307 
308 /*
309  *   This helper function is used by gen_close and gen_throw to
310  *   close a subiterator being delegated to by yield-from.
311  */
312 
313 static int
gen_close_iter(PyObject * yf)314 gen_close_iter(PyObject *yf)
315 {
316     PyObject *retval = NULL;
317     _Py_IDENTIFIER(close);
318 
319     if (PyGen_CheckExact(yf) || PyCoro_CheckExact(yf)) {
320         retval = gen_close((PyGenObject *)yf, NULL);
321         if (retval == NULL)
322             return -1;
323     }
324     else {
325         PyObject *meth;
326         if (_PyObject_LookupAttrId(yf, &PyId_close, &meth) < 0) {
327             PyErr_WriteUnraisable(yf);
328         }
329         if (meth) {
330             retval = _PyObject_CallNoArg(meth);
331             Py_DECREF(meth);
332             if (retval == NULL)
333                 return -1;
334         }
335     }
336     Py_XDECREF(retval);
337     return 0;
338 }
339 
340 PyObject *
_PyGen_yf(PyGenObject * gen)341 _PyGen_yf(PyGenObject *gen)
342 {
343     PyObject *yf = NULL;
344     PyFrameObject *f = gen->gi_frame;
345 
346     if (f) {
347         PyObject *bytecode = f->f_code->co_code;
348         unsigned char *code = (unsigned char *)PyBytes_AS_STRING(bytecode);
349 
350         if (f->f_lasti < 0) {
351             /* Return immediately if the frame didn't start yet. YIELD_FROM
352                always come after LOAD_CONST: a code object should not start
353                with YIELD_FROM */
354             assert(code[0] != YIELD_FROM);
355             return NULL;
356         }
357 
358         if (code[(f->f_lasti+1)*sizeof(_Py_CODEUNIT)] != YIELD_FROM)
359             return NULL;
360         assert(f->f_stackdepth > 0);
361         yf = f->f_valuestack[f->f_stackdepth-1];
362         Py_INCREF(yf);
363     }
364 
365     return yf;
366 }
367 
368 static PyObject *
gen_close(PyGenObject * gen,PyObject * args)369 gen_close(PyGenObject *gen, PyObject *args)
370 {
371     PyObject *retval;
372     PyObject *yf = _PyGen_yf(gen);
373     int err = 0;
374 
375     if (yf) {
376         PyFrameState state = gen->gi_frame->f_state;
377         gen->gi_frame->f_state = FRAME_EXECUTING;
378         err = gen_close_iter(yf);
379         gen->gi_frame->f_state = state;
380         Py_DECREF(yf);
381     }
382     if (err == 0)
383         PyErr_SetNone(PyExc_GeneratorExit);
384     retval = gen_send_ex(gen, Py_None, 1, 1);
385     if (retval) {
386         const char *msg = "generator ignored GeneratorExit";
387         if (PyCoro_CheckExact(gen)) {
388             msg = "coroutine ignored GeneratorExit";
389         } else if (PyAsyncGen_CheckExact(gen)) {
390             msg = ASYNC_GEN_IGNORED_EXIT_MSG;
391         }
392         Py_DECREF(retval);
393         PyErr_SetString(PyExc_RuntimeError, msg);
394         return NULL;
395     }
396     if (PyErr_ExceptionMatches(PyExc_StopIteration)
397         || PyErr_ExceptionMatches(PyExc_GeneratorExit)) {
398         PyErr_Clear();          /* ignore these errors */
399         Py_RETURN_NONE;
400     }
401     return NULL;
402 }
403 
404 
405 PyDoc_STRVAR(throw_doc,
406 "throw(typ[,val[,tb]]) -> raise exception in generator,\n\
407 return next yielded value or raise StopIteration.");
408 
409 static PyObject *
_gen_throw(PyGenObject * gen,int close_on_genexit,PyObject * typ,PyObject * val,PyObject * tb)410 _gen_throw(PyGenObject *gen, int close_on_genexit,
411            PyObject *typ, PyObject *val, PyObject *tb)
412 {
413     PyObject *yf = _PyGen_yf(gen);
414     _Py_IDENTIFIER(throw);
415 
416     if (yf) {
417         PyObject *ret;
418         int err;
419         if (PyErr_GivenExceptionMatches(typ, PyExc_GeneratorExit) &&
420             close_on_genexit
421         ) {
422             /* Asynchronous generators *should not* be closed right away.
423                We have to allow some awaits to work it through, hence the
424                `close_on_genexit` parameter here.
425             */
426             PyFrameState state = gen->gi_frame->f_state;
427             gen->gi_frame->f_state = FRAME_EXECUTING;
428             err = gen_close_iter(yf);
429             gen->gi_frame->f_state = state;
430             Py_DECREF(yf);
431             if (err < 0)
432                 return gen_send_ex(gen, Py_None, 1, 0);
433             goto throw_here;
434         }
435         if (PyGen_CheckExact(yf) || PyCoro_CheckExact(yf)) {
436             /* `yf` is a generator or a coroutine. */
437             PyThreadState *tstate = _PyThreadState_GET();
438             PyFrameObject *f = tstate->frame;
439 
440             /* Since we are fast-tracking things by skipping the eval loop,
441                we need to update the current frame so the stack trace
442                will be reported correctly to the user. */
443             /* XXX We should probably be updating the current frame
444                somewhere in ceval.c. */
445             tstate->frame = gen->gi_frame;
446             /* Close the generator that we are currently iterating with
447                'yield from' or awaiting on with 'await'. */
448             PyFrameState state = gen->gi_frame->f_state;
449             gen->gi_frame->f_state = FRAME_EXECUTING;
450             ret = _gen_throw((PyGenObject *)yf, close_on_genexit,
451                              typ, val, tb);
452             gen->gi_frame->f_state = state;
453             tstate->frame = f;
454         } else {
455             /* `yf` is an iterator or a coroutine-like object. */
456             PyObject *meth;
457             if (_PyObject_LookupAttrId(yf, &PyId_throw, &meth) < 0) {
458                 Py_DECREF(yf);
459                 return NULL;
460             }
461             if (meth == NULL) {
462                 Py_DECREF(yf);
463                 goto throw_here;
464             }
465             PyFrameState state = gen->gi_frame->f_state;
466             gen->gi_frame->f_state = FRAME_EXECUTING;
467             ret = PyObject_CallFunctionObjArgs(meth, typ, val, tb, NULL);
468             gen->gi_frame->f_state = state;
469             Py_DECREF(meth);
470         }
471         Py_DECREF(yf);
472         if (!ret) {
473             PyObject *val;
474             /* Pop subiterator from stack */
475             assert(gen->gi_frame->f_stackdepth > 0);
476             gen->gi_frame->f_stackdepth--;
477             ret = gen->gi_frame->f_valuestack[gen->gi_frame->f_stackdepth];
478             assert(ret == yf);
479             Py_DECREF(ret);
480             /* Termination repetition of YIELD_FROM */
481             assert(gen->gi_frame->f_lasti >= 0);
482             gen->gi_frame->f_lasti += 1;
483             if (_PyGen_FetchStopIterationValue(&val) == 0) {
484                 ret = gen_send(gen, val);
485                 Py_DECREF(val);
486             } else {
487                 ret = gen_send_ex(gen, Py_None, 1, 0);
488             }
489         }
490         return ret;
491     }
492 
493 throw_here:
494     /* First, check the traceback argument, replacing None with
495        NULL. */
496     if (tb == Py_None) {
497         tb = NULL;
498     }
499     else if (tb != NULL && !PyTraceBack_Check(tb)) {
500         PyErr_SetString(PyExc_TypeError,
501             "throw() third argument must be a traceback object");
502         return NULL;
503     }
504 
505     Py_INCREF(typ);
506     Py_XINCREF(val);
507     Py_XINCREF(tb);
508 
509     if (PyExceptionClass_Check(typ))
510         PyErr_NormalizeException(&typ, &val, &tb);
511 
512     else if (PyExceptionInstance_Check(typ)) {
513         /* Raising an instance.  The value should be a dummy. */
514         if (val && val != Py_None) {
515             PyErr_SetString(PyExc_TypeError,
516               "instance exception may not have a separate value");
517             goto failed_throw;
518         }
519         else {
520             /* Normalize to raise <class>, <instance> */
521             Py_XDECREF(val);
522             val = typ;
523             typ = PyExceptionInstance_Class(typ);
524             Py_INCREF(typ);
525 
526             if (tb == NULL)
527                 /* Returns NULL if there's no traceback */
528                 tb = PyException_GetTraceback(val);
529         }
530     }
531     else {
532         /* Not something you can raise.  throw() fails. */
533         PyErr_Format(PyExc_TypeError,
534                      "exceptions must be classes or instances "
535                      "deriving from BaseException, not %s",
536                      Py_TYPE(typ)->tp_name);
537             goto failed_throw;
538     }
539 
540     PyErr_Restore(typ, val, tb);
541     return gen_send_ex(gen, Py_None, 1, 0);
542 
543 failed_throw:
544     /* Didn't use our arguments, so restore their original refcounts */
545     Py_DECREF(typ);
546     Py_XDECREF(val);
547     Py_XDECREF(tb);
548     return NULL;
549 }
550 
551 
552 static PyObject *
gen_throw(PyGenObject * gen,PyObject * args)553 gen_throw(PyGenObject *gen, PyObject *args)
554 {
555     PyObject *typ;
556     PyObject *tb = NULL;
557     PyObject *val = NULL;
558 
559     if (!PyArg_UnpackTuple(args, "throw", 1, 3, &typ, &val, &tb)) {
560         return NULL;
561     }
562 
563     return _gen_throw(gen, 1, typ, val, tb);
564 }
565 
566 
567 static PyObject *
gen_iternext(PyGenObject * gen)568 gen_iternext(PyGenObject *gen)
569 {
570     PyObject *result;
571     assert(PyGen_CheckExact(gen) || PyCoro_CheckExact(gen));
572     if (gen_send_ex2(gen, NULL, &result, 0, 0) == PYGEN_RETURN) {
573         if (result != Py_None) {
574             _PyGen_SetStopIterationValue(result);
575         }
576         Py_CLEAR(result);
577     }
578     return result;
579 }
580 
581 /*
582  * Set StopIteration with specified value.  Value can be arbitrary object
583  * or NULL.
584  *
585  * Returns 0 if StopIteration is set and -1 if any other exception is set.
586  */
587 int
_PyGen_SetStopIterationValue(PyObject * value)588 _PyGen_SetStopIterationValue(PyObject *value)
589 {
590     PyObject *e;
591 
592     if (value == NULL ||
593         (!PyTuple_Check(value) && !PyExceptionInstance_Check(value)))
594     {
595         /* Delay exception instantiation if we can */
596         PyErr_SetObject(PyExc_StopIteration, value);
597         return 0;
598     }
599     /* Construct an exception instance manually with
600      * PyObject_CallOneArg and pass it to PyErr_SetObject.
601      *
602      * We do this to handle a situation when "value" is a tuple, in which
603      * case PyErr_SetObject would set the value of StopIteration to
604      * the first element of the tuple.
605      *
606      * (See PyErr_SetObject/_PyErr_CreateException code for details.)
607      */
608     e = PyObject_CallOneArg(PyExc_StopIteration, value);
609     if (e == NULL) {
610         return -1;
611     }
612     PyErr_SetObject(PyExc_StopIteration, e);
613     Py_DECREF(e);
614     return 0;
615 }
616 
617 /*
618  *   If StopIteration exception is set, fetches its 'value'
619  *   attribute if any, otherwise sets pvalue to None.
620  *
621  *   Returns 0 if no exception or StopIteration is set.
622  *   If any other exception is set, returns -1 and leaves
623  *   pvalue unchanged.
624  */
625 
626 int
_PyGen_FetchStopIterationValue(PyObject ** pvalue)627 _PyGen_FetchStopIterationValue(PyObject **pvalue)
628 {
629     PyObject *et, *ev, *tb;
630     PyObject *value = NULL;
631 
632     if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
633         PyErr_Fetch(&et, &ev, &tb);
634         if (ev) {
635             /* exception will usually be normalised already */
636             if (PyObject_TypeCheck(ev, (PyTypeObject *) et)) {
637                 value = ((PyStopIterationObject *)ev)->value;
638                 Py_INCREF(value);
639                 Py_DECREF(ev);
640             } else if (et == PyExc_StopIteration && !PyTuple_Check(ev)) {
641                 /* Avoid normalisation and take ev as value.
642                  *
643                  * Normalization is required if the value is a tuple, in
644                  * that case the value of StopIteration would be set to
645                  * the first element of the tuple.
646                  *
647                  * (See _PyErr_CreateException code for details.)
648                  */
649                 value = ev;
650             } else {
651                 /* normalisation required */
652                 PyErr_NormalizeException(&et, &ev, &tb);
653                 if (!PyObject_TypeCheck(ev, (PyTypeObject *)PyExc_StopIteration)) {
654                     PyErr_Restore(et, ev, tb);
655                     return -1;
656                 }
657                 value = ((PyStopIterationObject *)ev)->value;
658                 Py_INCREF(value);
659                 Py_DECREF(ev);
660             }
661         }
662         Py_XDECREF(et);
663         Py_XDECREF(tb);
664     } else if (PyErr_Occurred()) {
665         return -1;
666     }
667     if (value == NULL) {
668         value = Py_None;
669         Py_INCREF(value);
670     }
671     *pvalue = value;
672     return 0;
673 }
674 
675 static PyObject *
gen_repr(PyGenObject * gen)676 gen_repr(PyGenObject *gen)
677 {
678     return PyUnicode_FromFormat("<generator object %S at %p>",
679                                 gen->gi_qualname, gen);
680 }
681 
682 static PyObject *
gen_get_name(PyGenObject * op,void * Py_UNUSED (ignored))683 gen_get_name(PyGenObject *op, void *Py_UNUSED(ignored))
684 {
685     Py_INCREF(op->gi_name);
686     return op->gi_name;
687 }
688 
689 static int
gen_set_name(PyGenObject * op,PyObject * value,void * Py_UNUSED (ignored))690 gen_set_name(PyGenObject *op, PyObject *value, void *Py_UNUSED(ignored))
691 {
692     /* Not legal to del gen.gi_name or to set it to anything
693      * other than a string object. */
694     if (value == NULL || !PyUnicode_Check(value)) {
695         PyErr_SetString(PyExc_TypeError,
696                         "__name__ must be set to a string object");
697         return -1;
698     }
699     Py_INCREF(value);
700     Py_XSETREF(op->gi_name, value);
701     return 0;
702 }
703 
704 static PyObject *
gen_get_qualname(PyGenObject * op,void * Py_UNUSED (ignored))705 gen_get_qualname(PyGenObject *op, void *Py_UNUSED(ignored))
706 {
707     Py_INCREF(op->gi_qualname);
708     return op->gi_qualname;
709 }
710 
711 static int
gen_set_qualname(PyGenObject * op,PyObject * value,void * Py_UNUSED (ignored))712 gen_set_qualname(PyGenObject *op, PyObject *value, void *Py_UNUSED(ignored))
713 {
714     /* Not legal to del gen.__qualname__ or to set it to anything
715      * other than a string object. */
716     if (value == NULL || !PyUnicode_Check(value)) {
717         PyErr_SetString(PyExc_TypeError,
718                         "__qualname__ must be set to a string object");
719         return -1;
720     }
721     Py_INCREF(value);
722     Py_XSETREF(op->gi_qualname, value);
723     return 0;
724 }
725 
726 static PyObject *
gen_getyieldfrom(PyGenObject * gen,void * Py_UNUSED (ignored))727 gen_getyieldfrom(PyGenObject *gen, void *Py_UNUSED(ignored))
728 {
729     PyObject *yf = _PyGen_yf(gen);
730     if (yf == NULL)
731         Py_RETURN_NONE;
732     return yf;
733 }
734 
735 
736 static PyObject *
gen_getrunning(PyGenObject * gen,void * Py_UNUSED (ignored))737 gen_getrunning(PyGenObject *gen, void *Py_UNUSED(ignored))
738 {
739     if (gen->gi_frame == NULL) {
740         Py_RETURN_FALSE;
741     }
742     return PyBool_FromLong(_PyFrame_IsExecuting(gen->gi_frame));
743 }
744 
745 static PyGetSetDef gen_getsetlist[] = {
746     {"__name__", (getter)gen_get_name, (setter)gen_set_name,
747      PyDoc_STR("name of the generator")},
748     {"__qualname__", (getter)gen_get_qualname, (setter)gen_set_qualname,
749      PyDoc_STR("qualified name of the generator")},
750     {"gi_yieldfrom", (getter)gen_getyieldfrom, NULL,
751      PyDoc_STR("object being iterated by yield from, or None")},
752     {"gi_running", (getter)gen_getrunning, NULL, NULL},
753     {NULL} /* Sentinel */
754 };
755 
756 static PyMemberDef gen_memberlist[] = {
757     {"gi_frame",     T_OBJECT, offsetof(PyGenObject, gi_frame),    READONLY|PY_AUDIT_READ},
758     {"gi_code",      T_OBJECT, offsetof(PyGenObject, gi_code),     READONLY|PY_AUDIT_READ},
759     {NULL}      /* Sentinel */
760 };
761 
762 static PyMethodDef gen_methods[] = {
763     {"send",(PyCFunction)gen_send, METH_O, send_doc},
764     {"throw",(PyCFunction)gen_throw, METH_VARARGS, throw_doc},
765     {"close",(PyCFunction)gen_close, METH_NOARGS, close_doc},
766     {NULL, NULL}        /* Sentinel */
767 };
768 
769 static PyAsyncMethods gen_as_async = {
770     0,                                          /* am_await */
771     0,                                          /* am_aiter */
772     0,                                          /* am_anext */
773     (sendfunc)PyGen_am_send,                    /* am_send  */
774 };
775 
776 
777 PyTypeObject PyGen_Type = {
778     PyVarObject_HEAD_INIT(&PyType_Type, 0)
779     "generator",                                /* tp_name */
780     sizeof(PyGenObject),                        /* tp_basicsize */
781     0,                                          /* tp_itemsize */
782     /* methods */
783     (destructor)gen_dealloc,                    /* tp_dealloc */
784     0,                                          /* tp_vectorcall_offset */
785     0,                                          /* tp_getattr */
786     0,                                          /* tp_setattr */
787     &gen_as_async,                              /* tp_as_async */
788     (reprfunc)gen_repr,                         /* tp_repr */
789     0,                                          /* tp_as_number */
790     0,                                          /* tp_as_sequence */
791     0,                                          /* tp_as_mapping */
792     0,                                          /* tp_hash */
793     0,                                          /* tp_call */
794     0,                                          /* tp_str */
795     PyObject_GenericGetAttr,                    /* tp_getattro */
796     0,                                          /* tp_setattro */
797     0,                                          /* tp_as_buffer */
798     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
799     0,                                          /* tp_doc */
800     (traverseproc)gen_traverse,                 /* tp_traverse */
801     0,                                          /* tp_clear */
802     0,                                          /* tp_richcompare */
803     offsetof(PyGenObject, gi_weakreflist),      /* tp_weaklistoffset */
804     PyObject_SelfIter,                          /* tp_iter */
805     (iternextfunc)gen_iternext,                 /* tp_iternext */
806     gen_methods,                                /* tp_methods */
807     gen_memberlist,                             /* tp_members */
808     gen_getsetlist,                             /* tp_getset */
809     0,                                          /* tp_base */
810     0,                                          /* tp_dict */
811 
812     0,                                          /* tp_descr_get */
813     0,                                          /* tp_descr_set */
814     0,                                          /* tp_dictoffset */
815     0,                                          /* tp_init */
816     0,                                          /* tp_alloc */
817     0,                                          /* tp_new */
818     0,                                          /* tp_free */
819     0,                                          /* tp_is_gc */
820     0,                                          /* tp_bases */
821     0,                                          /* tp_mro */
822     0,                                          /* tp_cache */
823     0,                                          /* tp_subclasses */
824     0,                                          /* tp_weaklist */
825     0,                                          /* tp_del */
826     0,                                          /* tp_version_tag */
827     _PyGen_Finalize,                            /* tp_finalize */
828 };
829 
830 static PyObject *
gen_new_with_qualname(PyTypeObject * type,PyFrameObject * f,PyObject * name,PyObject * qualname)831 gen_new_with_qualname(PyTypeObject *type, PyFrameObject *f,
832                       PyObject *name, PyObject *qualname)
833 {
834     PyGenObject *gen = PyObject_GC_New(PyGenObject, type);
835     if (gen == NULL) {
836         Py_DECREF(f);
837         return NULL;
838     }
839     gen->gi_frame = f;
840     f->f_gen = (PyObject *) gen;
841     Py_INCREF(f->f_code);
842     gen->gi_code = (PyObject *)(f->f_code);
843     gen->gi_weakreflist = NULL;
844     gen->gi_exc_state.exc_type = NULL;
845     gen->gi_exc_state.exc_value = NULL;
846     gen->gi_exc_state.exc_traceback = NULL;
847     gen->gi_exc_state.previous_item = NULL;
848     if (name != NULL)
849         gen->gi_name = name;
850     else
851         gen->gi_name = ((PyCodeObject *)gen->gi_code)->co_name;
852     Py_INCREF(gen->gi_name);
853     if (qualname != NULL)
854         gen->gi_qualname = qualname;
855     else
856         gen->gi_qualname = gen->gi_name;
857     Py_INCREF(gen->gi_qualname);
858     _PyObject_GC_TRACK(gen);
859     return (PyObject *)gen;
860 }
861 
862 PyObject *
PyGen_NewWithQualName(PyFrameObject * f,PyObject * name,PyObject * qualname)863 PyGen_NewWithQualName(PyFrameObject *f, PyObject *name, PyObject *qualname)
864 {
865     return gen_new_with_qualname(&PyGen_Type, f, name, qualname);
866 }
867 
868 PyObject *
PyGen_New(PyFrameObject * f)869 PyGen_New(PyFrameObject *f)
870 {
871     return gen_new_with_qualname(&PyGen_Type, f, NULL, NULL);
872 }
873 
874 /* Coroutine Object */
875 
876 typedef struct {
877     PyObject_HEAD
878     PyCoroObject *cw_coroutine;
879 } PyCoroWrapper;
880 
881 static int
gen_is_coroutine(PyObject * o)882 gen_is_coroutine(PyObject *o)
883 {
884     if (PyGen_CheckExact(o)) {
885         PyCodeObject *code = (PyCodeObject *)((PyGenObject*)o)->gi_code;
886         if (code->co_flags & CO_ITERABLE_COROUTINE) {
887             return 1;
888         }
889     }
890     return 0;
891 }
892 
893 /*
894  *   This helper function returns an awaitable for `o`:
895  *     - `o` if `o` is a coroutine-object;
896  *     - `type(o)->tp_as_async->am_await(o)`
897  *
898  *   Raises a TypeError if it's not possible to return
899  *   an awaitable and returns NULL.
900  */
901 PyObject *
_PyCoro_GetAwaitableIter(PyObject * o)902 _PyCoro_GetAwaitableIter(PyObject *o)
903 {
904     unaryfunc getter = NULL;
905     PyTypeObject *ot;
906 
907     if (PyCoro_CheckExact(o) || gen_is_coroutine(o)) {
908         /* 'o' is a coroutine. */
909         Py_INCREF(o);
910         return o;
911     }
912 
913     ot = Py_TYPE(o);
914     if (ot->tp_as_async != NULL) {
915         getter = ot->tp_as_async->am_await;
916     }
917     if (getter != NULL) {
918         PyObject *res = (*getter)(o);
919         if (res != NULL) {
920             if (PyCoro_CheckExact(res) || gen_is_coroutine(res)) {
921                 /* __await__ must return an *iterator*, not
922                    a coroutine or another awaitable (see PEP 492) */
923                 PyErr_SetString(PyExc_TypeError,
924                                 "__await__() returned a coroutine");
925                 Py_CLEAR(res);
926             } else if (!PyIter_Check(res)) {
927                 PyErr_Format(PyExc_TypeError,
928                              "__await__() returned non-iterator "
929                              "of type '%.100s'",
930                              Py_TYPE(res)->tp_name);
931                 Py_CLEAR(res);
932             }
933         }
934         return res;
935     }
936 
937     PyErr_Format(PyExc_TypeError,
938                  "object %.100s can't be used in 'await' expression",
939                  ot->tp_name);
940     return NULL;
941 }
942 
943 static PyObject *
coro_repr(PyCoroObject * coro)944 coro_repr(PyCoroObject *coro)
945 {
946     return PyUnicode_FromFormat("<coroutine object %S at %p>",
947                                 coro->cr_qualname, coro);
948 }
949 
950 static PyObject *
coro_await(PyCoroObject * coro)951 coro_await(PyCoroObject *coro)
952 {
953     PyCoroWrapper *cw = PyObject_GC_New(PyCoroWrapper, &_PyCoroWrapper_Type);
954     if (cw == NULL) {
955         return NULL;
956     }
957     Py_INCREF(coro);
958     cw->cw_coroutine = coro;
959     _PyObject_GC_TRACK(cw);
960     return (PyObject *)cw;
961 }
962 
963 static PyObject *
coro_get_cr_await(PyCoroObject * coro,void * Py_UNUSED (ignored))964 coro_get_cr_await(PyCoroObject *coro, void *Py_UNUSED(ignored))
965 {
966     PyObject *yf = _PyGen_yf((PyGenObject *) coro);
967     if (yf == NULL)
968         Py_RETURN_NONE;
969     return yf;
970 }
971 
972 static PyObject *
cr_getrunning(PyCoroObject * coro,void * Py_UNUSED (ignored))973 cr_getrunning(PyCoroObject *coro, void *Py_UNUSED(ignored))
974 {
975     if (coro->cr_frame == NULL) {
976         Py_RETURN_FALSE;
977     }
978     return PyBool_FromLong(_PyFrame_IsExecuting(coro->cr_frame));
979 }
980 
981 static PyGetSetDef coro_getsetlist[] = {
982     {"__name__", (getter)gen_get_name, (setter)gen_set_name,
983      PyDoc_STR("name of the coroutine")},
984     {"__qualname__", (getter)gen_get_qualname, (setter)gen_set_qualname,
985      PyDoc_STR("qualified name of the coroutine")},
986     {"cr_await", (getter)coro_get_cr_await, NULL,
987      PyDoc_STR("object being awaited on, or None")},
988     {"cr_running", (getter)cr_getrunning, NULL, NULL},
989     {NULL} /* Sentinel */
990 };
991 
992 static PyMemberDef coro_memberlist[] = {
993     {"cr_frame",     T_OBJECT, offsetof(PyCoroObject, cr_frame),    READONLY|PY_AUDIT_READ},
994     {"cr_code",      T_OBJECT, offsetof(PyCoroObject, cr_code),     READONLY|PY_AUDIT_READ},
995     {"cr_origin",    T_OBJECT, offsetof(PyCoroObject, cr_origin),   READONLY},
996     {NULL}      /* Sentinel */
997 };
998 
999 PyDoc_STRVAR(coro_send_doc,
1000 "send(arg) -> send 'arg' into coroutine,\n\
1001 return next iterated value or raise StopIteration.");
1002 
1003 PyDoc_STRVAR(coro_throw_doc,
1004 "throw(typ[,val[,tb]]) -> raise exception in coroutine,\n\
1005 return next iterated value or raise StopIteration.");
1006 
1007 PyDoc_STRVAR(coro_close_doc,
1008 "close() -> raise GeneratorExit inside coroutine.");
1009 
1010 static PyMethodDef coro_methods[] = {
1011     {"send",(PyCFunction)gen_send, METH_O, coro_send_doc},
1012     {"throw",(PyCFunction)gen_throw, METH_VARARGS, coro_throw_doc},
1013     {"close",(PyCFunction)gen_close, METH_NOARGS, coro_close_doc},
1014     {NULL, NULL}        /* Sentinel */
1015 };
1016 
1017 static PyAsyncMethods coro_as_async = {
1018     (unaryfunc)coro_await,                      /* am_await */
1019     0,                                          /* am_aiter */
1020     0,                                          /* am_anext */
1021     (sendfunc)PyGen_am_send,                    /* am_send  */
1022 };
1023 
1024 PyTypeObject PyCoro_Type = {
1025     PyVarObject_HEAD_INIT(&PyType_Type, 0)
1026     "coroutine",                                /* tp_name */
1027     sizeof(PyCoroObject),                       /* tp_basicsize */
1028     0,                                          /* tp_itemsize */
1029     /* methods */
1030     (destructor)gen_dealloc,                    /* tp_dealloc */
1031     0,                                          /* tp_vectorcall_offset */
1032     0,                                          /* tp_getattr */
1033     0,                                          /* tp_setattr */
1034     &coro_as_async,                             /* tp_as_async */
1035     (reprfunc)coro_repr,                        /* tp_repr */
1036     0,                                          /* tp_as_number */
1037     0,                                          /* tp_as_sequence */
1038     0,                                          /* tp_as_mapping */
1039     0,                                          /* tp_hash */
1040     0,                                          /* tp_call */
1041     0,                                          /* tp_str */
1042     PyObject_GenericGetAttr,                    /* tp_getattro */
1043     0,                                          /* tp_setattro */
1044     0,                                          /* tp_as_buffer */
1045     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
1046     0,                                          /* tp_doc */
1047     (traverseproc)gen_traverse,                 /* tp_traverse */
1048     0,                                          /* tp_clear */
1049     0,                                          /* tp_richcompare */
1050     offsetof(PyCoroObject, cr_weakreflist),     /* tp_weaklistoffset */
1051     0,                                          /* tp_iter */
1052     0,                                          /* tp_iternext */
1053     coro_methods,                               /* tp_methods */
1054     coro_memberlist,                            /* tp_members */
1055     coro_getsetlist,                            /* tp_getset */
1056     0,                                          /* tp_base */
1057     0,                                          /* tp_dict */
1058     0,                                          /* tp_descr_get */
1059     0,                                          /* tp_descr_set */
1060     0,                                          /* tp_dictoffset */
1061     0,                                          /* tp_init */
1062     0,                                          /* tp_alloc */
1063     0,                                          /* tp_new */
1064     0,                                          /* tp_free */
1065     0,                                          /* tp_is_gc */
1066     0,                                          /* tp_bases */
1067     0,                                          /* tp_mro */
1068     0,                                          /* tp_cache */
1069     0,                                          /* tp_subclasses */
1070     0,                                          /* tp_weaklist */
1071     0,                                          /* tp_del */
1072     0,                                          /* tp_version_tag */
1073     _PyGen_Finalize,                            /* tp_finalize */
1074 };
1075 
1076 static void
coro_wrapper_dealloc(PyCoroWrapper * cw)1077 coro_wrapper_dealloc(PyCoroWrapper *cw)
1078 {
1079     _PyObject_GC_UNTRACK((PyObject *)cw);
1080     Py_CLEAR(cw->cw_coroutine);
1081     PyObject_GC_Del(cw);
1082 }
1083 
1084 static PyObject *
coro_wrapper_iternext(PyCoroWrapper * cw)1085 coro_wrapper_iternext(PyCoroWrapper *cw)
1086 {
1087     return gen_iternext((PyGenObject *)cw->cw_coroutine);
1088 }
1089 
1090 static PyObject *
coro_wrapper_send(PyCoroWrapper * cw,PyObject * arg)1091 coro_wrapper_send(PyCoroWrapper *cw, PyObject *arg)
1092 {
1093     return gen_send((PyGenObject *)cw->cw_coroutine, arg);
1094 }
1095 
1096 static PyObject *
coro_wrapper_throw(PyCoroWrapper * cw,PyObject * args)1097 coro_wrapper_throw(PyCoroWrapper *cw, PyObject *args)
1098 {
1099     return gen_throw((PyGenObject *)cw->cw_coroutine, args);
1100 }
1101 
1102 static PyObject *
coro_wrapper_close(PyCoroWrapper * cw,PyObject * args)1103 coro_wrapper_close(PyCoroWrapper *cw, PyObject *args)
1104 {
1105     return gen_close((PyGenObject *)cw->cw_coroutine, args);
1106 }
1107 
1108 static int
coro_wrapper_traverse(PyCoroWrapper * cw,visitproc visit,void * arg)1109 coro_wrapper_traverse(PyCoroWrapper *cw, visitproc visit, void *arg)
1110 {
1111     Py_VISIT((PyObject *)cw->cw_coroutine);
1112     return 0;
1113 }
1114 
1115 static PyMethodDef coro_wrapper_methods[] = {
1116     {"send",(PyCFunction)coro_wrapper_send, METH_O, coro_send_doc},
1117     {"throw",(PyCFunction)coro_wrapper_throw, METH_VARARGS, coro_throw_doc},
1118     {"close",(PyCFunction)coro_wrapper_close, METH_NOARGS, coro_close_doc},
1119     {NULL, NULL}        /* Sentinel */
1120 };
1121 
1122 PyTypeObject _PyCoroWrapper_Type = {
1123     PyVarObject_HEAD_INIT(&PyType_Type, 0)
1124     "coroutine_wrapper",
1125     sizeof(PyCoroWrapper),                      /* tp_basicsize */
1126     0,                                          /* tp_itemsize */
1127     (destructor)coro_wrapper_dealloc,           /* destructor tp_dealloc */
1128     0,                                          /* tp_vectorcall_offset */
1129     0,                                          /* tp_getattr */
1130     0,                                          /* tp_setattr */
1131     0,                                          /* tp_as_async */
1132     0,                                          /* tp_repr */
1133     0,                                          /* tp_as_number */
1134     0,                                          /* tp_as_sequence */
1135     0,                                          /* tp_as_mapping */
1136     0,                                          /* tp_hash */
1137     0,                                          /* tp_call */
1138     0,                                          /* tp_str */
1139     PyObject_GenericGetAttr,                    /* tp_getattro */
1140     0,                                          /* tp_setattro */
1141     0,                                          /* tp_as_buffer */
1142     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
1143     "A wrapper object implementing __await__ for coroutines.",
1144     (traverseproc)coro_wrapper_traverse,        /* tp_traverse */
1145     0,                                          /* tp_clear */
1146     0,                                          /* tp_richcompare */
1147     0,                                          /* tp_weaklistoffset */
1148     PyObject_SelfIter,                          /* tp_iter */
1149     (iternextfunc)coro_wrapper_iternext,        /* tp_iternext */
1150     coro_wrapper_methods,                       /* tp_methods */
1151     0,                                          /* tp_members */
1152     0,                                          /* tp_getset */
1153     0,                                          /* tp_base */
1154     0,                                          /* tp_dict */
1155     0,                                          /* tp_descr_get */
1156     0,                                          /* tp_descr_set */
1157     0,                                          /* tp_dictoffset */
1158     0,                                          /* tp_init */
1159     0,                                          /* tp_alloc */
1160     0,                                          /* tp_new */
1161     0,                                          /* tp_free */
1162 };
1163 
1164 static PyObject *
compute_cr_origin(int origin_depth)1165 compute_cr_origin(int origin_depth)
1166 {
1167     PyFrameObject *frame = PyEval_GetFrame();
1168     /* First count how many frames we have */
1169     int frame_count = 0;
1170     for (; frame && frame_count < origin_depth; ++frame_count) {
1171         frame = frame->f_back;
1172     }
1173 
1174     /* Now collect them */
1175     PyObject *cr_origin = PyTuple_New(frame_count);
1176     if (cr_origin == NULL) {
1177         return NULL;
1178     }
1179     frame = PyEval_GetFrame();
1180     for (int i = 0; i < frame_count; ++i) {
1181         PyCodeObject *code = frame->f_code;
1182         PyObject *frameinfo = Py_BuildValue("OiO",
1183                                             code->co_filename,
1184                                             PyFrame_GetLineNumber(frame),
1185                                             code->co_name);
1186         if (!frameinfo) {
1187             Py_DECREF(cr_origin);
1188             return NULL;
1189         }
1190         PyTuple_SET_ITEM(cr_origin, i, frameinfo);
1191         frame = frame->f_back;
1192     }
1193 
1194     return cr_origin;
1195 }
1196 
1197 PyObject *
PyCoro_New(PyFrameObject * f,PyObject * name,PyObject * qualname)1198 PyCoro_New(PyFrameObject *f, PyObject *name, PyObject *qualname)
1199 {
1200     PyObject *coro = gen_new_with_qualname(&PyCoro_Type, f, name, qualname);
1201     if (!coro) {
1202         return NULL;
1203     }
1204 
1205     PyThreadState *tstate = _PyThreadState_GET();
1206     int origin_depth = tstate->coroutine_origin_tracking_depth;
1207 
1208     if (origin_depth == 0) {
1209         ((PyCoroObject *)coro)->cr_origin = NULL;
1210     } else {
1211         PyObject *cr_origin = compute_cr_origin(origin_depth);
1212         ((PyCoroObject *)coro)->cr_origin = cr_origin;
1213         if (!cr_origin) {
1214             Py_DECREF(coro);
1215             return NULL;
1216         }
1217     }
1218 
1219     return coro;
1220 }
1221 
1222 
1223 /* ========= Asynchronous Generators ========= */
1224 
1225 
1226 typedef enum {
1227     AWAITABLE_STATE_INIT,   /* new awaitable, has not yet been iterated */
1228     AWAITABLE_STATE_ITER,   /* being iterated */
1229     AWAITABLE_STATE_CLOSED, /* closed */
1230 } AwaitableState;
1231 
1232 
1233 typedef struct PyAsyncGenASend {
1234     PyObject_HEAD
1235     PyAsyncGenObject *ags_gen;
1236 
1237     /* Can be NULL, when in the __anext__() mode
1238        (equivalent of "asend(None)") */
1239     PyObject *ags_sendval;
1240 
1241     AwaitableState ags_state;
1242 } PyAsyncGenASend;
1243 
1244 
1245 typedef struct PyAsyncGenAThrow {
1246     PyObject_HEAD
1247     PyAsyncGenObject *agt_gen;
1248 
1249     /* Can be NULL, when in the "aclose()" mode
1250        (equivalent of "athrow(GeneratorExit)") */
1251     PyObject *agt_args;
1252 
1253     AwaitableState agt_state;
1254 } PyAsyncGenAThrow;
1255 
1256 
1257 typedef struct _PyAsyncGenWrappedValue {
1258     PyObject_HEAD
1259     PyObject *agw_val;
1260 } _PyAsyncGenWrappedValue;
1261 
1262 
1263 #define _PyAsyncGenWrappedValue_CheckExact(o) \
1264                     Py_IS_TYPE(o, &_PyAsyncGenWrappedValue_Type)
1265 
1266 #define PyAsyncGenASend_CheckExact(o) \
1267                     Py_IS_TYPE(o, &_PyAsyncGenASend_Type)
1268 
1269 
1270 static int
async_gen_traverse(PyAsyncGenObject * gen,visitproc visit,void * arg)1271 async_gen_traverse(PyAsyncGenObject *gen, visitproc visit, void *arg)
1272 {
1273     Py_VISIT(gen->ag_finalizer);
1274     return gen_traverse((PyGenObject*)gen, visit, arg);
1275 }
1276 
1277 
1278 static PyObject *
async_gen_repr(PyAsyncGenObject * o)1279 async_gen_repr(PyAsyncGenObject *o)
1280 {
1281     return PyUnicode_FromFormat("<async_generator object %S at %p>",
1282                                 o->ag_qualname, o);
1283 }
1284 
1285 
1286 static int
async_gen_init_hooks(PyAsyncGenObject * o)1287 async_gen_init_hooks(PyAsyncGenObject *o)
1288 {
1289     PyThreadState *tstate;
1290     PyObject *finalizer;
1291     PyObject *firstiter;
1292 
1293     if (o->ag_hooks_inited) {
1294         return 0;
1295     }
1296 
1297     o->ag_hooks_inited = 1;
1298 
1299     tstate = _PyThreadState_GET();
1300 
1301     finalizer = tstate->async_gen_finalizer;
1302     if (finalizer) {
1303         Py_INCREF(finalizer);
1304         o->ag_finalizer = finalizer;
1305     }
1306 
1307     firstiter = tstate->async_gen_firstiter;
1308     if (firstiter) {
1309         PyObject *res;
1310 
1311         Py_INCREF(firstiter);
1312         res = PyObject_CallOneArg(firstiter, (PyObject *)o);
1313         Py_DECREF(firstiter);
1314         if (res == NULL) {
1315             return 1;
1316         }
1317         Py_DECREF(res);
1318     }
1319 
1320     return 0;
1321 }
1322 
1323 
1324 static PyObject *
async_gen_anext(PyAsyncGenObject * o)1325 async_gen_anext(PyAsyncGenObject *o)
1326 {
1327     if (async_gen_init_hooks(o)) {
1328         return NULL;
1329     }
1330     return async_gen_asend_new(o, NULL);
1331 }
1332 
1333 
1334 static PyObject *
async_gen_asend(PyAsyncGenObject * o,PyObject * arg)1335 async_gen_asend(PyAsyncGenObject *o, PyObject *arg)
1336 {
1337     if (async_gen_init_hooks(o)) {
1338         return NULL;
1339     }
1340     return async_gen_asend_new(o, arg);
1341 }
1342 
1343 
1344 static PyObject *
async_gen_aclose(PyAsyncGenObject * o,PyObject * arg)1345 async_gen_aclose(PyAsyncGenObject *o, PyObject *arg)
1346 {
1347     if (async_gen_init_hooks(o)) {
1348         return NULL;
1349     }
1350     return async_gen_athrow_new(o, NULL);
1351 }
1352 
1353 static PyObject *
async_gen_athrow(PyAsyncGenObject * o,PyObject * args)1354 async_gen_athrow(PyAsyncGenObject *o, PyObject *args)
1355 {
1356     if (async_gen_init_hooks(o)) {
1357         return NULL;
1358     }
1359     return async_gen_athrow_new(o, args);
1360 }
1361 
1362 
1363 static PyGetSetDef async_gen_getsetlist[] = {
1364     {"__name__", (getter)gen_get_name, (setter)gen_set_name,
1365      PyDoc_STR("name of the async generator")},
1366     {"__qualname__", (getter)gen_get_qualname, (setter)gen_set_qualname,
1367      PyDoc_STR("qualified name of the async generator")},
1368     {"ag_await", (getter)coro_get_cr_await, NULL,
1369      PyDoc_STR("object being awaited on, or None")},
1370     {NULL} /* Sentinel */
1371 };
1372 
1373 static PyMemberDef async_gen_memberlist[] = {
1374     {"ag_frame",   T_OBJECT, offsetof(PyAsyncGenObject, ag_frame),   READONLY|PY_AUDIT_READ},
1375     {"ag_running", T_BOOL,   offsetof(PyAsyncGenObject, ag_running_async),
1376         READONLY},
1377     {"ag_code",    T_OBJECT, offsetof(PyAsyncGenObject, ag_code),    READONLY|PY_AUDIT_READ},
1378     {NULL}      /* Sentinel */
1379 };
1380 
1381 PyDoc_STRVAR(async_aclose_doc,
1382 "aclose() -> raise GeneratorExit inside generator.");
1383 
1384 PyDoc_STRVAR(async_asend_doc,
1385 "asend(v) -> send 'v' in generator.");
1386 
1387 PyDoc_STRVAR(async_athrow_doc,
1388 "athrow(typ[,val[,tb]]) -> raise exception in generator.");
1389 
1390 static PyMethodDef async_gen_methods[] = {
1391     {"asend", (PyCFunction)async_gen_asend, METH_O, async_asend_doc},
1392     {"athrow",(PyCFunction)async_gen_athrow, METH_VARARGS, async_athrow_doc},
1393     {"aclose", (PyCFunction)async_gen_aclose, METH_NOARGS, async_aclose_doc},
1394     {"__class_getitem__",    (PyCFunction)Py_GenericAlias,
1395     METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},
1396     {NULL, NULL}        /* Sentinel */
1397 };
1398 
1399 
1400 static PyAsyncMethods async_gen_as_async = {
1401     0,                                          /* am_await */
1402     PyObject_SelfIter,                          /* am_aiter */
1403     (unaryfunc)async_gen_anext,                 /* am_anext */
1404     (sendfunc)PyGen_am_send,                    /* am_send  */
1405 };
1406 
1407 
1408 PyTypeObject PyAsyncGen_Type = {
1409     PyVarObject_HEAD_INIT(&PyType_Type, 0)
1410     "async_generator",                          /* tp_name */
1411     sizeof(PyAsyncGenObject),                   /* tp_basicsize */
1412     0,                                          /* tp_itemsize */
1413     /* methods */
1414     (destructor)gen_dealloc,                    /* tp_dealloc */
1415     0,                                          /* tp_vectorcall_offset */
1416     0,                                          /* tp_getattr */
1417     0,                                          /* tp_setattr */
1418     &async_gen_as_async,                        /* tp_as_async */
1419     (reprfunc)async_gen_repr,                   /* tp_repr */
1420     0,                                          /* tp_as_number */
1421     0,                                          /* tp_as_sequence */
1422     0,                                          /* tp_as_mapping */
1423     0,                                          /* tp_hash */
1424     0,                                          /* tp_call */
1425     0,                                          /* tp_str */
1426     PyObject_GenericGetAttr,                    /* tp_getattro */
1427     0,                                          /* tp_setattro */
1428     0,                                          /* tp_as_buffer */
1429     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
1430     0,                                          /* tp_doc */
1431     (traverseproc)async_gen_traverse,           /* tp_traverse */
1432     0,                                          /* tp_clear */
1433     0,                                          /* tp_richcompare */
1434     offsetof(PyAsyncGenObject, ag_weakreflist), /* tp_weaklistoffset */
1435     0,                                          /* tp_iter */
1436     0,                                          /* tp_iternext */
1437     async_gen_methods,                          /* tp_methods */
1438     async_gen_memberlist,                       /* tp_members */
1439     async_gen_getsetlist,                       /* tp_getset */
1440     0,                                          /* tp_base */
1441     0,                                          /* tp_dict */
1442     0,                                          /* tp_descr_get */
1443     0,                                          /* tp_descr_set */
1444     0,                                          /* tp_dictoffset */
1445     0,                                          /* tp_init */
1446     0,                                          /* tp_alloc */
1447     0,                                          /* tp_new */
1448     0,                                          /* tp_free */
1449     0,                                          /* tp_is_gc */
1450     0,                                          /* tp_bases */
1451     0,                                          /* tp_mro */
1452     0,                                          /* tp_cache */
1453     0,                                          /* tp_subclasses */
1454     0,                                          /* tp_weaklist */
1455     0,                                          /* tp_del */
1456     0,                                          /* tp_version_tag */
1457     _PyGen_Finalize,                            /* tp_finalize */
1458 };
1459 
1460 
1461 static struct _Py_async_gen_state *
get_async_gen_state(void)1462 get_async_gen_state(void)
1463 {
1464     PyInterpreterState *interp = _PyInterpreterState_GET();
1465     return &interp->async_gen;
1466 }
1467 
1468 
1469 PyObject *
PyAsyncGen_New(PyFrameObject * f,PyObject * name,PyObject * qualname)1470 PyAsyncGen_New(PyFrameObject *f, PyObject *name, PyObject *qualname)
1471 {
1472     PyAsyncGenObject *o;
1473     o = (PyAsyncGenObject *)gen_new_with_qualname(
1474         &PyAsyncGen_Type, f, name, qualname);
1475     if (o == NULL) {
1476         return NULL;
1477     }
1478     o->ag_finalizer = NULL;
1479     o->ag_closed = 0;
1480     o->ag_hooks_inited = 0;
1481     o->ag_running_async = 0;
1482     return (PyObject*)o;
1483 }
1484 
1485 
1486 void
_PyAsyncGen_ClearFreeLists(PyInterpreterState * interp)1487 _PyAsyncGen_ClearFreeLists(PyInterpreterState *interp)
1488 {
1489     struct _Py_async_gen_state *state = &interp->async_gen;
1490 
1491     while (state->value_numfree) {
1492         _PyAsyncGenWrappedValue *o;
1493         o = state->value_freelist[--state->value_numfree];
1494         assert(_PyAsyncGenWrappedValue_CheckExact(o));
1495         PyObject_GC_Del(o);
1496     }
1497 
1498     while (state->asend_numfree) {
1499         PyAsyncGenASend *o;
1500         o = state->asend_freelist[--state->asend_numfree];
1501         assert(Py_IS_TYPE(o, &_PyAsyncGenASend_Type));
1502         PyObject_GC_Del(o);
1503     }
1504 }
1505 
1506 void
_PyAsyncGen_Fini(PyInterpreterState * interp)1507 _PyAsyncGen_Fini(PyInterpreterState *interp)
1508 {
1509     _PyAsyncGen_ClearFreeLists(interp);
1510 #ifdef Py_DEBUG
1511     struct _Py_async_gen_state *state = &interp->async_gen;
1512     state->value_numfree = -1;
1513     state->asend_numfree = -1;
1514 #endif
1515 }
1516 
1517 
1518 static PyObject *
async_gen_unwrap_value(PyAsyncGenObject * gen,PyObject * result)1519 async_gen_unwrap_value(PyAsyncGenObject *gen, PyObject *result)
1520 {
1521     if (result == NULL) {
1522         if (!PyErr_Occurred()) {
1523             PyErr_SetNone(PyExc_StopAsyncIteration);
1524         }
1525 
1526         if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)
1527             || PyErr_ExceptionMatches(PyExc_GeneratorExit)
1528         ) {
1529             gen->ag_closed = 1;
1530         }
1531 
1532         gen->ag_running_async = 0;
1533         return NULL;
1534     }
1535 
1536     if (_PyAsyncGenWrappedValue_CheckExact(result)) {
1537         /* async yield */
1538         _PyGen_SetStopIterationValue(((_PyAsyncGenWrappedValue*)result)->agw_val);
1539         Py_DECREF(result);
1540         gen->ag_running_async = 0;
1541         return NULL;
1542     }
1543 
1544     return result;
1545 }
1546 
1547 
1548 /* ---------- Async Generator ASend Awaitable ------------ */
1549 
1550 
1551 static void
async_gen_asend_dealloc(PyAsyncGenASend * o)1552 async_gen_asend_dealloc(PyAsyncGenASend *o)
1553 {
1554     _PyObject_GC_UNTRACK((PyObject *)o);
1555     Py_CLEAR(o->ags_gen);
1556     Py_CLEAR(o->ags_sendval);
1557     struct _Py_async_gen_state *state = get_async_gen_state();
1558 #ifdef Py_DEBUG
1559     // async_gen_asend_dealloc() must not be called after _PyAsyncGen_Fini()
1560     assert(state->asend_numfree != -1);
1561 #endif
1562     if (state->asend_numfree < _PyAsyncGen_MAXFREELIST) {
1563         assert(PyAsyncGenASend_CheckExact(o));
1564         state->asend_freelist[state->asend_numfree++] = o;
1565     }
1566     else {
1567         PyObject_GC_Del(o);
1568     }
1569 }
1570 
1571 static int
async_gen_asend_traverse(PyAsyncGenASend * o,visitproc visit,void * arg)1572 async_gen_asend_traverse(PyAsyncGenASend *o, visitproc visit, void *arg)
1573 {
1574     Py_VISIT(o->ags_gen);
1575     Py_VISIT(o->ags_sendval);
1576     return 0;
1577 }
1578 
1579 
1580 static PyObject *
async_gen_asend_send(PyAsyncGenASend * o,PyObject * arg)1581 async_gen_asend_send(PyAsyncGenASend *o, PyObject *arg)
1582 {
1583     PyObject *result;
1584 
1585     if (o->ags_state == AWAITABLE_STATE_CLOSED) {
1586         PyErr_SetString(
1587             PyExc_RuntimeError,
1588             "cannot reuse already awaited __anext__()/asend()");
1589         return NULL;
1590     }
1591 
1592     if (o->ags_state == AWAITABLE_STATE_INIT) {
1593         if (o->ags_gen->ag_running_async) {
1594             PyErr_SetString(
1595                 PyExc_RuntimeError,
1596                 "anext(): asynchronous generator is already running");
1597             return NULL;
1598         }
1599 
1600         if (arg == NULL || arg == Py_None) {
1601             arg = o->ags_sendval;
1602         }
1603         o->ags_state = AWAITABLE_STATE_ITER;
1604     }
1605 
1606     o->ags_gen->ag_running_async = 1;
1607     result = gen_send((PyGenObject*)o->ags_gen, arg);
1608     result = async_gen_unwrap_value(o->ags_gen, result);
1609 
1610     if (result == NULL) {
1611         o->ags_state = AWAITABLE_STATE_CLOSED;
1612     }
1613 
1614     return result;
1615 }
1616 
1617 
1618 static PyObject *
async_gen_asend_iternext(PyAsyncGenASend * o)1619 async_gen_asend_iternext(PyAsyncGenASend *o)
1620 {
1621     return async_gen_asend_send(o, NULL);
1622 }
1623 
1624 
1625 static PyObject *
async_gen_asend_throw(PyAsyncGenASend * o,PyObject * args)1626 async_gen_asend_throw(PyAsyncGenASend *o, PyObject *args)
1627 {
1628     PyObject *result;
1629 
1630     if (o->ags_state == AWAITABLE_STATE_CLOSED) {
1631         PyErr_SetString(
1632             PyExc_RuntimeError,
1633             "cannot reuse already awaited __anext__()/asend()");
1634         return NULL;
1635     }
1636 
1637     result = gen_throw((PyGenObject*)o->ags_gen, args);
1638     result = async_gen_unwrap_value(o->ags_gen, result);
1639 
1640     if (result == NULL) {
1641         o->ags_state = AWAITABLE_STATE_CLOSED;
1642     }
1643 
1644     return result;
1645 }
1646 
1647 
1648 static PyObject *
async_gen_asend_close(PyAsyncGenASend * o,PyObject * args)1649 async_gen_asend_close(PyAsyncGenASend *o, PyObject *args)
1650 {
1651     o->ags_state = AWAITABLE_STATE_CLOSED;
1652     Py_RETURN_NONE;
1653 }
1654 
1655 
1656 static PyMethodDef async_gen_asend_methods[] = {
1657     {"send", (PyCFunction)async_gen_asend_send, METH_O, send_doc},
1658     {"throw", (PyCFunction)async_gen_asend_throw, METH_VARARGS, throw_doc},
1659     {"close", (PyCFunction)async_gen_asend_close, METH_NOARGS, close_doc},
1660     {NULL, NULL}        /* Sentinel */
1661 };
1662 
1663 
1664 static PyAsyncMethods async_gen_asend_as_async = {
1665     PyObject_SelfIter,                          /* am_await */
1666     0,                                          /* am_aiter */
1667     0,                                          /* am_anext */
1668     0,                                          /* am_send  */
1669 };
1670 
1671 
1672 PyTypeObject _PyAsyncGenASend_Type = {
1673     PyVarObject_HEAD_INIT(&PyType_Type, 0)
1674     "async_generator_asend",                    /* tp_name */
1675     sizeof(PyAsyncGenASend),                    /* tp_basicsize */
1676     0,                                          /* tp_itemsize */
1677     /* methods */
1678     (destructor)async_gen_asend_dealloc,        /* tp_dealloc */
1679     0,                                          /* tp_vectorcall_offset */
1680     0,                                          /* tp_getattr */
1681     0,                                          /* tp_setattr */
1682     &async_gen_asend_as_async,                  /* tp_as_async */
1683     0,                                          /* tp_repr */
1684     0,                                          /* tp_as_number */
1685     0,                                          /* tp_as_sequence */
1686     0,                                          /* tp_as_mapping */
1687     0,                                          /* tp_hash */
1688     0,                                          /* tp_call */
1689     0,                                          /* tp_str */
1690     PyObject_GenericGetAttr,                    /* tp_getattro */
1691     0,                                          /* tp_setattro */
1692     0,                                          /* tp_as_buffer */
1693     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
1694     0,                                          /* tp_doc */
1695     (traverseproc)async_gen_asend_traverse,     /* tp_traverse */
1696     0,                                          /* tp_clear */
1697     0,                                          /* tp_richcompare */
1698     0,                                          /* tp_weaklistoffset */
1699     PyObject_SelfIter,                          /* tp_iter */
1700     (iternextfunc)async_gen_asend_iternext,     /* tp_iternext */
1701     async_gen_asend_methods,                    /* tp_methods */
1702     0,                                          /* tp_members */
1703     0,                                          /* tp_getset */
1704     0,                                          /* tp_base */
1705     0,                                          /* tp_dict */
1706     0,                                          /* tp_descr_get */
1707     0,                                          /* tp_descr_set */
1708     0,                                          /* tp_dictoffset */
1709     0,                                          /* tp_init */
1710     0,                                          /* tp_alloc */
1711     0,                                          /* tp_new */
1712 };
1713 
1714 
1715 static PyObject *
async_gen_asend_new(PyAsyncGenObject * gen,PyObject * sendval)1716 async_gen_asend_new(PyAsyncGenObject *gen, PyObject *sendval)
1717 {
1718     PyAsyncGenASend *o;
1719     struct _Py_async_gen_state *state = get_async_gen_state();
1720 #ifdef Py_DEBUG
1721     // async_gen_asend_new() must not be called after _PyAsyncGen_Fini()
1722     assert(state->asend_numfree != -1);
1723 #endif
1724     if (state->asend_numfree) {
1725         state->asend_numfree--;
1726         o = state->asend_freelist[state->asend_numfree];
1727         _Py_NewReference((PyObject *)o);
1728     }
1729     else {
1730         o = PyObject_GC_New(PyAsyncGenASend, &_PyAsyncGenASend_Type);
1731         if (o == NULL) {
1732             return NULL;
1733         }
1734     }
1735 
1736     Py_INCREF(gen);
1737     o->ags_gen = gen;
1738 
1739     Py_XINCREF(sendval);
1740     o->ags_sendval = sendval;
1741 
1742     o->ags_state = AWAITABLE_STATE_INIT;
1743 
1744     _PyObject_GC_TRACK((PyObject*)o);
1745     return (PyObject*)o;
1746 }
1747 
1748 
1749 /* ---------- Async Generator Value Wrapper ------------ */
1750 
1751 
1752 static void
async_gen_wrapped_val_dealloc(_PyAsyncGenWrappedValue * o)1753 async_gen_wrapped_val_dealloc(_PyAsyncGenWrappedValue *o)
1754 {
1755     _PyObject_GC_UNTRACK((PyObject *)o);
1756     Py_CLEAR(o->agw_val);
1757     struct _Py_async_gen_state *state = get_async_gen_state();
1758 #ifdef Py_DEBUG
1759     // async_gen_wrapped_val_dealloc() must not be called after _PyAsyncGen_Fini()
1760     assert(state->value_numfree != -1);
1761 #endif
1762     if (state->value_numfree < _PyAsyncGen_MAXFREELIST) {
1763         assert(_PyAsyncGenWrappedValue_CheckExact(o));
1764         state->value_freelist[state->value_numfree++] = o;
1765     }
1766     else {
1767         PyObject_GC_Del(o);
1768     }
1769 }
1770 
1771 
1772 static int
async_gen_wrapped_val_traverse(_PyAsyncGenWrappedValue * o,visitproc visit,void * arg)1773 async_gen_wrapped_val_traverse(_PyAsyncGenWrappedValue *o,
1774                                visitproc visit, void *arg)
1775 {
1776     Py_VISIT(o->agw_val);
1777     return 0;
1778 }
1779 
1780 
1781 PyTypeObject _PyAsyncGenWrappedValue_Type = {
1782     PyVarObject_HEAD_INIT(&PyType_Type, 0)
1783     "async_generator_wrapped_value",            /* tp_name */
1784     sizeof(_PyAsyncGenWrappedValue),            /* tp_basicsize */
1785     0,                                          /* tp_itemsize */
1786     /* methods */
1787     (destructor)async_gen_wrapped_val_dealloc,  /* tp_dealloc */
1788     0,                                          /* tp_vectorcall_offset */
1789     0,                                          /* tp_getattr */
1790     0,                                          /* tp_setattr */
1791     0,                                          /* tp_as_async */
1792     0,                                          /* tp_repr */
1793     0,                                          /* tp_as_number */
1794     0,                                          /* tp_as_sequence */
1795     0,                                          /* tp_as_mapping */
1796     0,                                          /* tp_hash */
1797     0,                                          /* tp_call */
1798     0,                                          /* tp_str */
1799     PyObject_GenericGetAttr,                    /* tp_getattro */
1800     0,                                          /* tp_setattro */
1801     0,                                          /* tp_as_buffer */
1802     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
1803     0,                                          /* tp_doc */
1804     (traverseproc)async_gen_wrapped_val_traverse, /* tp_traverse */
1805     0,                                          /* tp_clear */
1806     0,                                          /* tp_richcompare */
1807     0,                                          /* tp_weaklistoffset */
1808     0,                                          /* tp_iter */
1809     0,                                          /* tp_iternext */
1810     0,                                          /* tp_methods */
1811     0,                                          /* tp_members */
1812     0,                                          /* tp_getset */
1813     0,                                          /* tp_base */
1814     0,                                          /* tp_dict */
1815     0,                                          /* tp_descr_get */
1816     0,                                          /* tp_descr_set */
1817     0,                                          /* tp_dictoffset */
1818     0,                                          /* tp_init */
1819     0,                                          /* tp_alloc */
1820     0,                                          /* tp_new */
1821 };
1822 
1823 
1824 PyObject *
_PyAsyncGenValueWrapperNew(PyObject * val)1825 _PyAsyncGenValueWrapperNew(PyObject *val)
1826 {
1827     _PyAsyncGenWrappedValue *o;
1828     assert(val);
1829 
1830     struct _Py_async_gen_state *state = get_async_gen_state();
1831 #ifdef Py_DEBUG
1832     // _PyAsyncGenValueWrapperNew() must not be called after _PyAsyncGen_Fini()
1833     assert(state->value_numfree != -1);
1834 #endif
1835     if (state->value_numfree) {
1836         state->value_numfree--;
1837         o = state->value_freelist[state->value_numfree];
1838         assert(_PyAsyncGenWrappedValue_CheckExact(o));
1839         _Py_NewReference((PyObject*)o);
1840     }
1841     else {
1842         o = PyObject_GC_New(_PyAsyncGenWrappedValue,
1843                             &_PyAsyncGenWrappedValue_Type);
1844         if (o == NULL) {
1845             return NULL;
1846         }
1847     }
1848     o->agw_val = val;
1849     Py_INCREF(val);
1850     _PyObject_GC_TRACK((PyObject*)o);
1851     return (PyObject*)o;
1852 }
1853 
1854 
1855 /* ---------- Async Generator AThrow awaitable ------------ */
1856 
1857 
1858 static void
async_gen_athrow_dealloc(PyAsyncGenAThrow * o)1859 async_gen_athrow_dealloc(PyAsyncGenAThrow *o)
1860 {
1861     _PyObject_GC_UNTRACK((PyObject *)o);
1862     Py_CLEAR(o->agt_gen);
1863     Py_CLEAR(o->agt_args);
1864     PyObject_GC_Del(o);
1865 }
1866 
1867 
1868 static int
async_gen_athrow_traverse(PyAsyncGenAThrow * o,visitproc visit,void * arg)1869 async_gen_athrow_traverse(PyAsyncGenAThrow *o, visitproc visit, void *arg)
1870 {
1871     Py_VISIT(o->agt_gen);
1872     Py_VISIT(o->agt_args);
1873     return 0;
1874 }
1875 
1876 
1877 static PyObject *
async_gen_athrow_send(PyAsyncGenAThrow * o,PyObject * arg)1878 async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
1879 {
1880     PyGenObject *gen = (PyGenObject*)o->agt_gen;
1881     PyFrameObject *f = gen->gi_frame;
1882     PyObject *retval;
1883 
1884     if (o->agt_state == AWAITABLE_STATE_CLOSED) {
1885         PyErr_SetString(
1886             PyExc_RuntimeError,
1887             "cannot reuse already awaited aclose()/athrow()");
1888         return NULL;
1889     }
1890 
1891     if (f == NULL || _PyFrameHasCompleted(f)) {
1892         o->agt_state = AWAITABLE_STATE_CLOSED;
1893         PyErr_SetNone(PyExc_StopIteration);
1894         return NULL;
1895     }
1896 
1897     if (o->agt_state == AWAITABLE_STATE_INIT) {
1898         if (o->agt_gen->ag_running_async) {
1899             o->agt_state = AWAITABLE_STATE_CLOSED;
1900             if (o->agt_args == NULL) {
1901                 PyErr_SetString(
1902                     PyExc_RuntimeError,
1903                     "aclose(): asynchronous generator is already running");
1904             }
1905             else {
1906                 PyErr_SetString(
1907                     PyExc_RuntimeError,
1908                     "athrow(): asynchronous generator is already running");
1909             }
1910             return NULL;
1911         }
1912 
1913         if (o->agt_gen->ag_closed) {
1914             o->agt_state = AWAITABLE_STATE_CLOSED;
1915             PyErr_SetNone(PyExc_StopAsyncIteration);
1916             return NULL;
1917         }
1918 
1919         if (arg != Py_None) {
1920             PyErr_SetString(PyExc_RuntimeError, NON_INIT_CORO_MSG);
1921             return NULL;
1922         }
1923 
1924         o->agt_state = AWAITABLE_STATE_ITER;
1925         o->agt_gen->ag_running_async = 1;
1926 
1927         if (o->agt_args == NULL) {
1928             /* aclose() mode */
1929             o->agt_gen->ag_closed = 1;
1930 
1931             retval = _gen_throw((PyGenObject *)gen,
1932                                 0,  /* Do not close generator when
1933                                        PyExc_GeneratorExit is passed */
1934                                 PyExc_GeneratorExit, NULL, NULL);
1935 
1936             if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) {
1937                 Py_DECREF(retval);
1938                 goto yield_close;
1939             }
1940         } else {
1941             PyObject *typ;
1942             PyObject *tb = NULL;
1943             PyObject *val = NULL;
1944 
1945             if (!PyArg_UnpackTuple(o->agt_args, "athrow", 1, 3,
1946                                    &typ, &val, &tb)) {
1947                 return NULL;
1948             }
1949 
1950             retval = _gen_throw((PyGenObject *)gen,
1951                                 0,  /* Do not close generator when
1952                                        PyExc_GeneratorExit is passed */
1953                                 typ, val, tb);
1954             retval = async_gen_unwrap_value(o->agt_gen, retval);
1955         }
1956         if (retval == NULL) {
1957             goto check_error;
1958         }
1959         return retval;
1960     }
1961 
1962     assert(o->agt_state == AWAITABLE_STATE_ITER);
1963 
1964     retval = gen_send((PyGenObject *)gen, arg);
1965     if (o->agt_args) {
1966         return async_gen_unwrap_value(o->agt_gen, retval);
1967     } else {
1968         /* aclose() mode */
1969         if (retval) {
1970             if (_PyAsyncGenWrappedValue_CheckExact(retval)) {
1971                 Py_DECREF(retval);
1972                 goto yield_close;
1973             }
1974             else {
1975                 return retval;
1976             }
1977         }
1978         else {
1979             goto check_error;
1980         }
1981     }
1982 
1983 yield_close:
1984     o->agt_gen->ag_running_async = 0;
1985     o->agt_state = AWAITABLE_STATE_CLOSED;
1986     PyErr_SetString(
1987         PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
1988     return NULL;
1989 
1990 check_error:
1991     o->agt_gen->ag_running_async = 0;
1992     o->agt_state = AWAITABLE_STATE_CLOSED;
1993     if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) ||
1994             PyErr_ExceptionMatches(PyExc_GeneratorExit))
1995     {
1996         if (o->agt_args == NULL) {
1997             /* when aclose() is called we don't want to propagate
1998                StopAsyncIteration or GeneratorExit; just raise
1999                StopIteration, signalling that this 'aclose()' await
2000                is done.
2001             */
2002             PyErr_Clear();
2003             PyErr_SetNone(PyExc_StopIteration);
2004         }
2005     }
2006     return NULL;
2007 }
2008 
2009 
2010 static PyObject *
async_gen_athrow_throw(PyAsyncGenAThrow * o,PyObject * args)2011 async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *args)
2012 {
2013     PyObject *retval;
2014 
2015     if (o->agt_state == AWAITABLE_STATE_CLOSED) {
2016         PyErr_SetString(
2017             PyExc_RuntimeError,
2018             "cannot reuse already awaited aclose()/athrow()");
2019         return NULL;
2020     }
2021 
2022     retval = gen_throw((PyGenObject*)o->agt_gen, args);
2023     if (o->agt_args) {
2024         return async_gen_unwrap_value(o->agt_gen, retval);
2025     } else {
2026         /* aclose() mode */
2027         if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) {
2028             o->agt_gen->ag_running_async = 0;
2029             o->agt_state = AWAITABLE_STATE_CLOSED;
2030             Py_DECREF(retval);
2031             PyErr_SetString(PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
2032             return NULL;
2033         }
2034         if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) ||
2035             PyErr_ExceptionMatches(PyExc_GeneratorExit))
2036         {
2037             /* when aclose() is called we don't want to propagate
2038                StopAsyncIteration or GeneratorExit; just raise
2039                StopIteration, signalling that this 'aclose()' await
2040                is done.
2041             */
2042             PyErr_Clear();
2043             PyErr_SetNone(PyExc_StopIteration);
2044         }
2045         return retval;
2046     }
2047 }
2048 
2049 
2050 static PyObject *
async_gen_athrow_iternext(PyAsyncGenAThrow * o)2051 async_gen_athrow_iternext(PyAsyncGenAThrow *o)
2052 {
2053     return async_gen_athrow_send(o, Py_None);
2054 }
2055 
2056 
2057 static PyObject *
async_gen_athrow_close(PyAsyncGenAThrow * o,PyObject * args)2058 async_gen_athrow_close(PyAsyncGenAThrow *o, PyObject *args)
2059 {
2060     o->agt_state = AWAITABLE_STATE_CLOSED;
2061     Py_RETURN_NONE;
2062 }
2063 
2064 
2065 static PyMethodDef async_gen_athrow_methods[] = {
2066     {"send", (PyCFunction)async_gen_athrow_send, METH_O, send_doc},
2067     {"throw", (PyCFunction)async_gen_athrow_throw, METH_VARARGS, throw_doc},
2068     {"close", (PyCFunction)async_gen_athrow_close, METH_NOARGS, close_doc},
2069     {NULL, NULL}        /* Sentinel */
2070 };
2071 
2072 
2073 static PyAsyncMethods async_gen_athrow_as_async = {
2074     PyObject_SelfIter,                          /* am_await */
2075     0,                                          /* am_aiter */
2076     0,                                          /* am_anext */
2077     0,                                          /* am_send  */
2078 };
2079 
2080 
2081 PyTypeObject _PyAsyncGenAThrow_Type = {
2082     PyVarObject_HEAD_INIT(&PyType_Type, 0)
2083     "async_generator_athrow",                   /* tp_name */
2084     sizeof(PyAsyncGenAThrow),                   /* tp_basicsize */
2085     0,                                          /* tp_itemsize */
2086     /* methods */
2087     (destructor)async_gen_athrow_dealloc,       /* tp_dealloc */
2088     0,                                          /* tp_vectorcall_offset */
2089     0,                                          /* tp_getattr */
2090     0,                                          /* tp_setattr */
2091     &async_gen_athrow_as_async,                 /* tp_as_async */
2092     0,                                          /* tp_repr */
2093     0,                                          /* tp_as_number */
2094     0,                                          /* tp_as_sequence */
2095     0,                                          /* tp_as_mapping */
2096     0,                                          /* tp_hash */
2097     0,                                          /* tp_call */
2098     0,                                          /* tp_str */
2099     PyObject_GenericGetAttr,                    /* tp_getattro */
2100     0,                                          /* tp_setattro */
2101     0,                                          /* tp_as_buffer */
2102     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
2103     0,                                          /* tp_doc */
2104     (traverseproc)async_gen_athrow_traverse,    /* tp_traverse */
2105     0,                                          /* tp_clear */
2106     0,                                          /* tp_richcompare */
2107     0,                                          /* tp_weaklistoffset */
2108     PyObject_SelfIter,                          /* tp_iter */
2109     (iternextfunc)async_gen_athrow_iternext,    /* tp_iternext */
2110     async_gen_athrow_methods,                   /* tp_methods */
2111     0,                                          /* tp_members */
2112     0,                                          /* tp_getset */
2113     0,                                          /* tp_base */
2114     0,                                          /* tp_dict */
2115     0,                                          /* tp_descr_get */
2116     0,                                          /* tp_descr_set */
2117     0,                                          /* tp_dictoffset */
2118     0,                                          /* tp_init */
2119     0,                                          /* tp_alloc */
2120     0,                                          /* tp_new */
2121 };
2122 
2123 
2124 static PyObject *
async_gen_athrow_new(PyAsyncGenObject * gen,PyObject * args)2125 async_gen_athrow_new(PyAsyncGenObject *gen, PyObject *args)
2126 {
2127     PyAsyncGenAThrow *o;
2128     o = PyObject_GC_New(PyAsyncGenAThrow, &_PyAsyncGenAThrow_Type);
2129     if (o == NULL) {
2130         return NULL;
2131     }
2132     o->agt_gen = gen;
2133     o->agt_args = args;
2134     o->agt_state = AWAITABLE_STATE_INIT;
2135     Py_INCREF(gen);
2136     Py_XINCREF(args);
2137     _PyObject_GC_TRACK((PyObject*)o);
2138     return (PyObject*)o;
2139 }
2140