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