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