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