1 /* Generator object implementation */
2 
3 #include "Python.h"
4 #include "frameobject.h"
5 #include "genobject.h"
6 #include "ceval.h"
7 #include "structmember.h"
8 #include "opcode.h"
9 
10 static int
gen_traverse(PyGenObject * gen,visitproc visit,void * arg)11 gen_traverse(PyGenObject *gen, visitproc visit, void *arg)
12 {
13     Py_VISIT((PyObject *)gen->gi_frame);
14     Py_VISIT(gen->gi_code);
15     return 0;
16 }
17 
18 static void
gen_dealloc(PyGenObject * gen)19 gen_dealloc(PyGenObject *gen)
20 {
21     PyObject *self = (PyObject *) gen;
22 
23     _PyObject_GC_UNTRACK(gen);
24 
25     if (gen->gi_weakreflist != NULL)
26         PyObject_ClearWeakRefs(self);
27 
28     _PyObject_GC_TRACK(self);
29 
30     if (gen->gi_frame != NULL && gen->gi_frame->f_stacktop != NULL) {
31         /* Generator is paused, so we need to close */
32         Py_TYPE(gen)->tp_del(self);
33         if (self->ob_refcnt > 0)
34             return;                     /* resurrected.  :( */
35     }
36 
37     _PyObject_GC_UNTRACK(self);
38     Py_CLEAR(gen->gi_frame);
39     Py_CLEAR(gen->gi_code);
40     PyObject_GC_Del(gen);
41 }
42 
43 
44 static PyObject *
gen_send_ex(PyGenObject * gen,PyObject * arg,int exc)45 gen_send_ex(PyGenObject *gen, PyObject *arg, int exc)
46 {
47     PyThreadState *tstate = PyThreadState_GET();
48     PyFrameObject *f = gen->gi_frame;
49     PyObject *result;
50 
51     if (gen->gi_running) {
52         PyErr_SetString(PyExc_ValueError,
53                         "generator already executing");
54         return NULL;
55     }
56     if (f==NULL || f->f_stacktop == NULL) {
57         /* Only set exception if called from send() */
58         if (arg && !exc)
59             PyErr_SetNone(PyExc_StopIteration);
60         return NULL;
61     }
62 
63     if (f->f_lasti == -1) {
64         if (arg && arg != Py_None) {
65             PyErr_SetString(PyExc_TypeError,
66                             "can't send non-None value to a "
67                             "just-started generator");
68             return NULL;
69         }
70     } else {
71         /* Push arg onto the frame's value stack */
72         result = arg ? arg : Py_None;
73         Py_INCREF(result);
74         *(f->f_stacktop++) = result;
75     }
76 
77     /* Generators always return to their most recent caller, not
78      * necessarily their creator. */
79     f->f_tstate = tstate;
80     Py_XINCREF(tstate->frame);
81     assert(f->f_back == NULL);
82     f->f_back = tstate->frame;
83 
84     gen->gi_running = 1;
85     result = PyEval_EvalFrameEx(f, exc);
86     gen->gi_running = 0;
87 
88     /* Don't keep the reference to f_back any longer than necessary.  It
89      * may keep a chain of frames alive or it could create a reference
90      * cycle. */
91     assert(f->f_back == tstate->frame);
92     Py_CLEAR(f->f_back);
93     /* Clear the borrowed reference to the thread state */
94     f->f_tstate = NULL;
95 
96     /* If the generator just returned (as opposed to yielding), signal
97      * that the generator is exhausted. */
98     if (result == Py_None && f->f_stacktop == NULL) {
99         Py_DECREF(result);
100         result = NULL;
101         /* Set exception if not called by gen_iternext() */
102         if (arg)
103             PyErr_SetNone(PyExc_StopIteration);
104     }
105 
106     if (!result || f->f_stacktop == NULL) {
107         /* generator can't be rerun, so release the frame */
108         Py_DECREF(f);
109         gen->gi_frame = NULL;
110     }
111 
112     return result;
113 }
114 
115 PyDoc_STRVAR(send_doc,
116 "send(arg) -> send 'arg' into generator,\n\
117 return next yielded value or raise StopIteration.");
118 
119 static PyObject *
gen_send(PyGenObject * gen,PyObject * arg)120 gen_send(PyGenObject *gen, PyObject *arg)
121 {
122     return gen_send_ex(gen, arg, 0);
123 }
124 
125 PyDoc_STRVAR(close_doc,
126 "close() -> raise GeneratorExit inside generator.");
127 
128 static PyObject *
gen_close(PyGenObject * gen,PyObject * args)129 gen_close(PyGenObject *gen, PyObject *args)
130 {
131     PyObject *retval;
132     PyErr_SetNone(PyExc_GeneratorExit);
133     retval = gen_send_ex(gen, Py_None, 1);
134     if (retval) {
135         Py_DECREF(retval);
136         PyErr_SetString(PyExc_RuntimeError,
137                         "generator ignored GeneratorExit");
138         return NULL;
139     }
140     if (PyErr_ExceptionMatches(PyExc_StopIteration)
141         || PyErr_ExceptionMatches(PyExc_GeneratorExit))
142     {
143         PyErr_Clear();          /* ignore these errors */
144         Py_INCREF(Py_None);
145         return Py_None;
146     }
147     return NULL;
148 }
149 
150 static void
gen_del(PyObject * self)151 gen_del(PyObject *self)
152 {
153     PyObject *res;
154     PyObject *error_type, *error_value, *error_traceback;
155     PyGenObject *gen = (PyGenObject *)self;
156 
157     if (gen->gi_frame == NULL || gen->gi_frame->f_stacktop == NULL)
158         /* Generator isn't paused, so no need to close */
159         return;
160 
161     /* Temporarily resurrect the object. */
162     assert(self->ob_refcnt == 0);
163     self->ob_refcnt = 1;
164 
165     /* Save the current exception, if any. */
166     PyErr_Fetch(&error_type, &error_value, &error_traceback);
167 
168     res = gen_close(gen, NULL);
169 
170     if (res == NULL)
171         PyErr_WriteUnraisable(self);
172     else
173         Py_DECREF(res);
174 
175     /* Restore the saved exception. */
176     PyErr_Restore(error_type, error_value, error_traceback);
177 
178     /* Undo the temporary resurrection; can't use DECREF here, it would
179      * cause a recursive call.
180      */
181     assert(self->ob_refcnt > 0);
182     if (--self->ob_refcnt == 0)
183         return; /* this is the normal path out */
184 
185     /* close() resurrected it!  Make it look like the original Py_DECREF
186      * never happened.
187      */
188     {
189         Py_ssize_t refcnt = self->ob_refcnt;
190         _Py_NewReference(self);
191         self->ob_refcnt = refcnt;
192     }
193     assert(PyType_IS_GC(self->ob_type) &&
194            _Py_AS_GC(self)->gc.gc_refs != _PyGC_REFS_UNTRACKED);
195 
196     /* If Py_REF_DEBUG, _Py_NewReference bumped _Py_RefTotal, so
197      * we need to undo that. */
198     _Py_DEC_REFTOTAL;
199     /* If Py_TRACE_REFS, _Py_NewReference re-added self to the object
200      * chain, so no more to do there.
201      * If COUNT_ALLOCS, the original decref bumped tp_frees, and
202      * _Py_NewReference bumped tp_allocs:  both of those need to be
203      * undone.
204      */
205 #ifdef COUNT_ALLOCS
206     --self->ob_type->tp_frees;
207     --self->ob_type->tp_allocs;
208 #endif
209 }
210 
211 
212 
213 PyDoc_STRVAR(throw_doc,
214 "throw(typ[,val[,tb]]) -> raise exception in generator,\n\
215 return next yielded value or raise StopIteration.");
216 
217 static PyObject *
gen_throw(PyGenObject * gen,PyObject * args)218 gen_throw(PyGenObject *gen, PyObject *args)
219 {
220     PyObject *typ;
221     PyObject *tb = NULL;
222     PyObject *val = NULL;
223 
224     if (!PyArg_UnpackTuple(args, "throw", 1, 3, &typ, &val, &tb))
225         return NULL;
226 
227     /* First, check the traceback argument, replacing None with
228        NULL. */
229     if (tb == Py_None)
230         tb = NULL;
231     else if (tb != NULL && !PyTraceBack_Check(tb)) {
232         PyErr_SetString(PyExc_TypeError,
233             "throw() third argument must be a traceback object");
234         return NULL;
235     }
236 
237     Py_INCREF(typ);
238     Py_XINCREF(val);
239     Py_XINCREF(tb);
240 
241     if (PyExceptionClass_Check(typ)) {
242         PyErr_NormalizeException(&typ, &val, &tb);
243     }
244 
245     else if (PyExceptionInstance_Check(typ)) {
246         /* Raising an instance.  The value should be a dummy. */
247         if (val && val != Py_None) {
248             PyErr_SetString(PyExc_TypeError,
249               "instance exception may not have a separate value");
250             goto failed_throw;
251         }
252         else {
253             /* Normalize to raise <class>, <instance> */
254             Py_XDECREF(val);
255             val = typ;
256             typ = PyExceptionInstance_Class(typ);
257             Py_INCREF(typ);
258         }
259     }
260     else {
261         /* Not something you can raise.  throw() fails. */
262         PyErr_Format(PyExc_TypeError,
263                      "exceptions must be classes, or instances, not %s",
264                      typ->ob_type->tp_name);
265             goto failed_throw;
266     }
267 
268     PyErr_Restore(typ, val, tb);
269     return gen_send_ex(gen, Py_None, 1);
270 
271 failed_throw:
272     /* Didn't use our arguments, so restore their original refcounts */
273     Py_DECREF(typ);
274     Py_XDECREF(val);
275     Py_XDECREF(tb);
276     return NULL;
277 }
278 
279 
280 static PyObject *
gen_iternext(PyGenObject * gen)281 gen_iternext(PyGenObject *gen)
282 {
283     return gen_send_ex(gen, NULL, 0);
284 }
285 
286 
287 static PyObject *
gen_repr(PyGenObject * gen)288 gen_repr(PyGenObject *gen)
289 {
290     char *code_name;
291     code_name = PyString_AsString(((PyCodeObject *)gen->gi_code)->co_name);
292     if (code_name == NULL)
293         return NULL;
294     return PyString_FromFormat("<generator object %.200s at %p>",
295                                code_name, gen);
296 }
297 
298 
299 static PyObject *
gen_get_name(PyGenObject * gen)300 gen_get_name(PyGenObject *gen)
301 {
302     PyObject *name = ((PyCodeObject *)gen->gi_code)->co_name;
303     Py_INCREF(name);
304     return name;
305 }
306 
307 
308 PyDoc_STRVAR(gen__name__doc__,
309 "Return the name of the generator's associated code object.");
310 
311 static PyGetSetDef gen_getsetlist[] = {
312     {"__name__", (getter)gen_get_name, NULL, gen__name__doc__},
313     {NULL}
314 };
315 
316 
317 static PyMemberDef gen_memberlist[] = {
318     {"gi_frame",        T_OBJECT, offsetof(PyGenObject, gi_frame),      RO},
319     {"gi_running",      T_INT,    offsetof(PyGenObject, gi_running),    RO},
320     {"gi_code",     T_OBJECT, offsetof(PyGenObject, gi_code),  RO},
321     {NULL}      /* Sentinel */
322 };
323 
324 static PyMethodDef gen_methods[] = {
325     {"send",(PyCFunction)gen_send, METH_O, send_doc},
326     {"throw",(PyCFunction)gen_throw, METH_VARARGS, throw_doc},
327     {"close",(PyCFunction)gen_close, METH_NOARGS, close_doc},
328     {NULL, NULL}        /* Sentinel */
329 };
330 
331 PyTypeObject PyGen_Type = {
332     PyVarObject_HEAD_INIT(&PyType_Type, 0)
333     "generator",                                /* tp_name */
334     sizeof(PyGenObject),                        /* tp_basicsize */
335     0,                                          /* tp_itemsize */
336     /* methods */
337     (destructor)gen_dealloc,                    /* tp_dealloc */
338     0,                                          /* tp_print */
339     0,                                          /* tp_getattr */
340     0,                                          /* tp_setattr */
341     0,                                          /* tp_compare */
342     (reprfunc)gen_repr,                         /* tp_repr */
343     0,                                          /* tp_as_number */
344     0,                                          /* tp_as_sequence */
345     0,                                          /* tp_as_mapping */
346     0,                                          /* tp_hash */
347     0,                                          /* tp_call */
348     0,                                          /* tp_str */
349     PyObject_GenericGetAttr,                    /* tp_getattro */
350     0,                                          /* tp_setattro */
351     0,                                          /* tp_as_buffer */
352     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
353     0,                                          /* tp_doc */
354     (traverseproc)gen_traverse,                 /* tp_traverse */
355     0,                                          /* tp_clear */
356     0,                                          /* tp_richcompare */
357     offsetof(PyGenObject, gi_weakreflist),      /* tp_weaklistoffset */
358     PyObject_SelfIter,                          /* tp_iter */
359     (iternextfunc)gen_iternext,                 /* tp_iternext */
360     gen_methods,                                /* tp_methods */
361     gen_memberlist,                             /* tp_members */
362     gen_getsetlist,                             /* tp_getset */
363     0,                                          /* tp_base */
364     0,                                          /* tp_dict */
365 
366     0,                                          /* tp_descr_get */
367     0,                                          /* tp_descr_set */
368     0,                                          /* tp_dictoffset */
369     0,                                          /* tp_init */
370     0,                                          /* tp_alloc */
371     0,                                          /* tp_new */
372     0,                                          /* tp_free */
373     0,                                          /* tp_is_gc */
374     0,                                          /* tp_bases */
375     0,                                          /* tp_mro */
376     0,                                          /* tp_cache */
377     0,                                          /* tp_subclasses */
378     0,                                          /* tp_weaklist */
379     gen_del,                                    /* tp_del */
380 };
381 
382 PyObject *
PyGen_New(PyFrameObject * f)383 PyGen_New(PyFrameObject *f)
384 {
385     PyGenObject *gen = PyObject_GC_New(PyGenObject, &PyGen_Type);
386     if (gen == NULL) {
387         Py_DECREF(f);
388         return NULL;
389     }
390     gen->gi_frame = f;
391     Py_INCREF(f->f_code);
392     gen->gi_code = (PyObject *)(f->f_code);
393     gen->gi_running = 0;
394     gen->gi_weakreflist = NULL;
395     _PyObject_GC_TRACK(gen);
396     return (PyObject *)gen;
397 }
398 
399 int
PyGen_NeedsFinalizing(PyGenObject * gen)400 PyGen_NeedsFinalizing(PyGenObject *gen)
401 {
402     int i;
403     PyFrameObject *f = gen->gi_frame;
404 
405     if (f == NULL || f->f_stacktop == NULL || f->f_iblock <= 0)
406         return 0; /* no frame or empty blockstack == no finalization */
407 
408     /* Any block type besides a loop requires cleanup. */
409     i = f->f_iblock;
410     while (--i >= 0) {
411         if (f->f_blockstack[i].b_type != SETUP_LOOP)
412             return 1;
413     }
414 
415     /* No blocks except loops, it's safe to skip finalization. */
416     return 0;
417 }
418