Objects/genobject.c

   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
  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
  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 *
  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         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
 112 PyDoc_STRVAR(send_doc,
 113 "send(arg) -> send 'arg' into generator,\n\
 114 return next yielded value or raise StopIteration.");
 115
 116 static PyObject *
 117 gen_send(PyGenObject *gen, PyObject *arg)
 118 {
 119         return gen_send_ex(gen, arg, 0);
 120 }
 121
 122 PyDoc_STRVAR(close_doc,
 123 "close(arg) -> raise GeneratorExit inside generator.");
 124
 125 static PyObject *
 126 gen_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
 147 static void
 148 gen_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
 210 PyDoc_STRVAR(throw_doc,
 211 "throw(typ[,val[,tb]]) -> raise exception in generator,\n\
 212 return next yielded value or raise StopIteration.");
 213
 214 static PyObject *
 215 gen_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
 268 failed_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
 277 static PyObject *
 278 gen_iternext(PyGenObject *gen)
 279 {
 280         return gen_send_ex(gen, NULL, 0);
 281 }
 282
 283
 284 static PyMemberDef gen_memberlist[] = {
 285         {"gi_frame",    T_OBJECT, offsetof(PyGenObject, gi_frame),      RO},
 286         {"gi_running",  T_INT,    offsetof(PyGenObject, gi_running),    RO},
 287         {"gi_code",     T_OBJECT, offsetof(PyGenObject, gi_code),  RO},
 288         {NULL}  /* Sentinel */
 289 };
 290
 291 static PyMethodDef gen_methods[] = {
 292         {"send",(PyCFunction)gen_send, METH_O, send_doc},
 293         {"throw",(PyCFunction)gen_throw, METH_VARARGS, throw_doc},
 294         {"close",(PyCFunction)gen_close, METH_NOARGS, close_doc},
 295         {NULL, NULL}    /* Sentinel */
 296 };
 297
 298 PyTypeObject PyGen_Type = {
 299         PyVarObject_HEAD_INIT(&PyType_Type, 0)
 300         "generator",                            /* tp_name */
 301         sizeof(PyGenObject),                    /* tp_basicsize */
 302         0,                                      /* tp_itemsize */
 303         /* methods */
 304         (destructor)gen_dealloc,                /* tp_dealloc */
 305         0,                                      /* tp_print */
 306         0,                                      /* tp_getattr */
 307         0,                                      /* tp_setattr */
 308         0,                                      /* tp_compare */
 309         0,                                      /* tp_repr */
 310         0,                                      /* tp_as_number */
 311         0,                                      /* tp_as_sequence */
 312         0,                                      /* tp_as_mapping */
 313         0,                                      /* tp_hash */
 314         0,                                      /* tp_call */
 315         0,                                      /* tp_str */
 316         PyObject_GenericGetAttr,                /* tp_getattro */
 317         0,                                      /* tp_setattro */
 318         0,                                      /* tp_as_buffer */
 319         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
 320         0,                                      /* tp_doc */
 321         (traverseproc)gen_traverse,             /* tp_traverse */
 322         0,                                      /* tp_clear */
 323         0,                                      /* tp_richcompare */
 324         offsetof(PyGenObject, gi_weakreflist),  /* tp_weaklistoffset */
 325         PyObject_SelfIter,                      /* tp_iter */
 326         (iternextfunc)gen_iternext,             /* tp_iternext */
 327         gen_methods,                            /* tp_methods */
 328         gen_memberlist,                         /* tp_members */
 329         0,                                      /* tp_getset */
 330         0,                                      /* tp_base */
 331         0,                                      /* tp_dict */
 332
 333         0,                                      /* tp_descr_get */
 334         0,                                      /* tp_descr_set */
 335         0,                                      /* tp_dictoffset */
 336         0,                                      /* tp_init */
 337         0,                                      /* tp_alloc */
 338         0,                                      /* tp_new */
 339         0,                                      /* tp_free */
 340         0,                                      /* tp_is_gc */
 341         0,                                      /* tp_bases */
 342         0,                                      /* tp_mro */
 343         0,                                      /* tp_cache */
 344         0,                                      /* tp_subclasses */
 345         0,                                      /* tp_weaklist */
 346         gen_del,                                /* tp_del */
 347 };
 348
 349 PyObject *
 350 PyGen_New(PyFrameObject *f)
 351 {
 352         PyGenObject *gen = PyObject_GC_New(PyGenObject, &PyGen_Type);
 353         if (gen == NULL) {
 354                 Py_DECREF(f);
 355                 return NULL;
 356         }
 357         gen->gi_frame = f;
 358         Py_INCREF(f->f_code);
 359         gen->gi_code = (PyObject *)(f->f_code);
 360         gen->gi_running = 0;
 361         gen->gi_weakreflist = NULL;
 362         _PyObject_GC_TRACK(gen);
 363         return (PyObject *)gen;
 364 }
 365
 366 int
 367 PyGen_NeedsFinalizing(PyGenObject *gen)
 368 {
 369         int i;
 370         PyFrameObject *f = gen->gi_frame;
 371
 372         if (f == NULL || f->f_stacktop == NULL || f->f_iblock <= 0)
 373                 return 0; /* no frame or empty blockstack == no finalization */
 374
 375         /* Any block type besides a loop requires cleanup. */
 376         i = f->f_iblock;
 377         while (--i >= 0) {
 378                 if (f->f_blockstack[i].b_type != SETUP_LOOP)
 379                         return 1;
 380         }
 381
 382         /* No blocks except loops, it's safe to skip finalization. */
 383         return 0;
 384 }