Objects/enumobject.c

   1 /* enumerate object */
   2
   3 #include "Python.h"
   4
   5 typedef struct {
   6     PyObject_HEAD
   7     Py_ssize_t en_index;           /* current index of enumeration */
   8     PyObject* en_sit;          /* secondary iterator of enumeration */
   9     PyObject* en_result;           /* result tuple  */
  10     PyObject* en_longindex;        /* index for sequences >= PY_SSIZE_T_MAX */
  11 } enumobject;
  12
  13 static PyObject *
  14 enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  15 {
  16     enumobject *en;
  17     PyObject *seq = NULL;
  18     PyObject *start = NULL;
  19     static char *kwlist[] = {"sequence", "start", 0};
  20
  21     if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:enumerate", kwlist,
  22                                      &seq, &start))
  23         return NULL;
  24
  25     en = (enumobject *)type->tp_alloc(type, 0);
  26     if (en == NULL)
  27         return NULL;
  28     if (start != NULL) {
  29         start = PyNumber_Index(start);
  30         if (start == NULL) {
  31             Py_DECREF(en);
  32             return NULL;
  33         }
  34         assert(PyInt_Check(start) || PyLong_Check(start));
  35         en->en_index = PyInt_AsSsize_t(start);
  36         if (en->en_index == -1 && PyErr_Occurred()) {
  37             PyErr_Clear();
  38             en->en_index = PY_SSIZE_T_MAX;
  39             en->en_longindex = start;
  40         } else {
  41             en->en_longindex = NULL;
  42             Py_DECREF(start);
  43         }
  44     } else {
  45         en->en_index = 0;
  46         en->en_longindex = NULL;
  47     }
  48     en->en_sit = PyObject_GetIter(seq);
  49     if (en->en_sit == NULL) {
  50         Py_DECREF(en);
  51         return NULL;
  52     }
  53     en->en_result = PyTuple_Pack(2, Py_None, Py_None);
  54     if (en->en_result == NULL) {
  55         Py_DECREF(en);
  56         return NULL;
  57     }
  58     return (PyObject *)en;
  59 }
  60
  61 static void
  62 enum_dealloc(enumobject *en)
  63 {
  64     PyObject_GC_UnTrack(en);
  65     Py_XDECREF(en->en_sit);
  66     Py_XDECREF(en->en_result);
  67     Py_XDECREF(en->en_longindex);
  68     Py_TYPE(en)->tp_free(en);
  69 }
  70
  71 static int
  72 enum_traverse(enumobject *en, visitproc visit, void *arg)
  73 {
  74     Py_VISIT(en->en_sit);
  75     Py_VISIT(en->en_result);
  76     Py_VISIT(en->en_longindex);
  77     return 0;
  78 }
  79
  80 static PyObject *
  81 enum_next_long(enumobject *en, PyObject* next_item)
  82 {
  83     static PyObject *one = NULL;
  84     PyObject *result = en->en_result;
  85     PyObject *next_index;
  86     PyObject *stepped_up;
  87
  88     if (en->en_longindex == NULL) {
  89         en->en_longindex = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
  90         if (en->en_longindex == NULL)
  91             return NULL;
  92     }
  93     if (one == NULL) {
  94         one = PyInt_FromLong(1);
  95         if (one == NULL)
  96             return NULL;
  97     }
  98     next_index = en->en_longindex;
  99     assert(next_index != NULL);
 100     stepped_up = PyNumber_Add(next_index, one);
 101     if (stepped_up == NULL)
 102         return NULL;
 103     en->en_longindex = stepped_up;
 104
 105     if (result->ob_refcnt == 1) {
 106         Py_INCREF(result);
 107         Py_DECREF(PyTuple_GET_ITEM(result, 0));
 108         Py_DECREF(PyTuple_GET_ITEM(result, 1));
 109     } else {
 110         result = PyTuple_New(2);
 111         if (result == NULL) {
 112             Py_DECREF(next_index);
 113             Py_DECREF(next_item);
 114             return NULL;
 115         }
 116     }
 117     PyTuple_SET_ITEM(result, 0, next_index);
 118     PyTuple_SET_ITEM(result, 1, next_item);
 119     return result;
 120 }
 121
 122 static PyObject *
 123 enum_next(enumobject *en)
 124 {
 125     PyObject *next_index;
 126     PyObject *next_item;
 127     PyObject *result = en->en_result;
 128     PyObject *it = en->en_sit;
 129
 130     next_item = (*Py_TYPE(it)->tp_iternext)(it);
 131     if (next_item == NULL)
 132         return NULL;
 133
 134     if (en->en_index == PY_SSIZE_T_MAX)
 135         return enum_next_long(en, next_item);
 136
 137     next_index = PyInt_FromSsize_t(en->en_index);
 138     if (next_index == NULL) {
 139         Py_DECREF(next_item);
 140         return NULL;
 141     }
 142     en->en_index++;
 143
 144     if (result->ob_refcnt == 1) {
 145         Py_INCREF(result);
 146         Py_DECREF(PyTuple_GET_ITEM(result, 0));
 147         Py_DECREF(PyTuple_GET_ITEM(result, 1));
 148     } else {
 149         result = PyTuple_New(2);
 150         if (result == NULL) {
 151             Py_DECREF(next_index);
 152             Py_DECREF(next_item);
 153             return NULL;
 154         }
 155     }
 156     PyTuple_SET_ITEM(result, 0, next_index);
 157     PyTuple_SET_ITEM(result, 1, next_item);
 158     return result;
 159 }
 160
 161 PyDoc_STRVAR(enum_doc,
 162 "enumerate(iterable[, start]) -> iterator for index, value of iterable\n"
 163 "\n"
 164 "Return an enumerate object.  iterable must be another object that supports\n"
 165 "iteration.  The enumerate object yields pairs containing a count (from\n"
 166 "start, which defaults to zero) and a value yielded by the iterable argument.\n"
 167 "enumerate is useful for obtaining an indexed list:\n"
 168 "    (0, seq[0]), (1, seq[1]), (2, seq[2]), ...");
 169
 170 PyTypeObject PyEnum_Type = {
 171     PyVarObject_HEAD_INIT(&PyType_Type, 0)
 172     "enumerate",                    /* tp_name */
 173     sizeof(enumobject),             /* tp_basicsize */
 174     0,                              /* tp_itemsize */
 175     /* methods */
 176     (destructor)enum_dealloc,       /* tp_dealloc */
 177     0,                              /* tp_print */
 178     0,                              /* tp_getattr */
 179     0,                              /* tp_setattr */
 180     0,                              /* tp_compare */
 181     0,                              /* tp_repr */
 182     0,                              /* tp_as_number */
 183     0,                              /* tp_as_sequence */
 184     0,                              /* tp_as_mapping */
 185     0,                              /* tp_hash */
 186     0,                              /* tp_call */
 187     0,                              /* tp_str */
 188     PyObject_GenericGetAttr,        /* tp_getattro */
 189     0,                              /* tp_setattro */
 190     0,                              /* tp_as_buffer */
 191     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 192         Py_TPFLAGS_BASETYPE,    /* tp_flags */
 193     enum_doc,                       /* tp_doc */
 194     (traverseproc)enum_traverse,    /* tp_traverse */
 195     0,                              /* tp_clear */
 196     0,                              /* tp_richcompare */
 197     0,                              /* tp_weaklistoffset */
 198     PyObject_SelfIter,                  /* tp_iter */
 199     (iternextfunc)enum_next,        /* tp_iternext */
 200     0,                              /* tp_methods */
 201     0,                              /* tp_members */
 202     0,                              /* tp_getset */
 203     0,                              /* tp_base */
 204     0,                              /* tp_dict */
 205     0,                              /* tp_descr_get */
 206     0,                              /* tp_descr_set */
 207     0,                              /* tp_dictoffset */
 208     0,                              /* tp_init */
 209     PyType_GenericAlloc,            /* tp_alloc */
 210     enum_new,                       /* tp_new */
 211     PyObject_GC_Del,                /* tp_free */
 212 };
 213
 214 /* Reversed Object ***************************************************************/
 215
 216 typedef struct {
 217     PyObject_HEAD
 218     Py_ssize_t      index;
 219     PyObject* seq;
 220 } reversedobject;
 221
 222 static PyObject *
 223 reversed_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 224 {
 225     Py_ssize_t n;
 226     PyObject *seq, *reversed_meth;
 227     static PyObject *reversed_cache = NULL;
 228     reversedobject *ro;
 229
 230     if (type == &PyReversed_Type && !_PyArg_NoKeywords("reversed()", kwds))
 231         return NULL;
 232
 233     if (!PyArg_UnpackTuple(args, "reversed", 1, 1, &seq) )
 234         return NULL;
 235
 236     if (PyInstance_Check(seq)) {
 237         reversed_meth = PyObject_GetAttrString(seq, "__reversed__");
 238         if (reversed_meth == NULL) {
 239             if (PyErr_ExceptionMatches(PyExc_AttributeError))
 240                 PyErr_Clear();
 241             else
 242                 return NULL;
 243         }
 244     }
 245     else {
 246         reversed_meth = _PyObject_LookupSpecial(seq, "__reversed__",
 247                                                 &reversed_cache);
 248         if (reversed_meth == NULL && PyErr_Occurred())
 249             return NULL;
 250     }
 251     if (reversed_meth != NULL) {
 252         PyObject *res = PyObject_CallFunctionObjArgs(reversed_meth, NULL);
 253         Py_DECREF(reversed_meth);
 254         return res;
 255     }
 256
 257     if (!PySequence_Check(seq)) {
 258         PyErr_SetString(PyExc_TypeError,
 259                         "argument to reversed() must be a sequence");
 260         return NULL;
 261     }
 262
 263     n = PySequence_Size(seq);
 264     if (n == -1)
 265         return NULL;
 266
 267     ro = (reversedobject *)type->tp_alloc(type, 0);
 268     if (ro == NULL)
 269         return NULL;
 270
 271     ro->index = n-1;
 272     Py_INCREF(seq);
 273     ro->seq = seq;
 274     return (PyObject *)ro;
 275 }
 276
 277 static void
 278 reversed_dealloc(reversedobject *ro)
 279 {
 280     PyObject_GC_UnTrack(ro);
 281     Py_XDECREF(ro->seq);
 282     Py_TYPE(ro)->tp_free(ro);
 283 }
 284
 285 static int
 286 reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
 287 {
 288     Py_VISIT(ro->seq);
 289     return 0;
 290 }
 291
 292 static PyObject *
 293 reversed_next(reversedobject *ro)
 294 {
 295     PyObject *item;
 296     Py_ssize_t index = ro->index;
 297
 298     if (index >= 0) {
 299         item = PySequence_GetItem(ro->seq, index);
 300         if (item != NULL) {
 301             ro->index--;
 302             return item;
 303         }
 304         if (PyErr_ExceptionMatches(PyExc_IndexError) ||
 305             PyErr_ExceptionMatches(PyExc_StopIteration))
 306             PyErr_Clear();
 307     }
 308     ro->index = -1;
 309     Py_CLEAR(ro->seq);
 310     return NULL;
 311 }
 312
 313 PyDoc_STRVAR(reversed_doc,
 314 "reversed(sequence) -> reverse iterator over values of the sequence\n"
 315 "\n"
 316 "Return a reverse iterator");
 317
 318 static PyObject *
 319 reversed_len(reversedobject *ro)
 320 {
 321     Py_ssize_t position, seqsize;
 322
 323     if (ro->seq == NULL)
 324         return PyInt_FromLong(0);
 325     seqsize = PySequence_Size(ro->seq);
 326     if (seqsize == -1)
 327         return NULL;
 328     position = ro->index + 1;
 329     return PyInt_FromSsize_t((seqsize < position)  ?  0  :  position);
 330 }
 331
 332 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
 333
 334 static PyMethodDef reversediter_methods[] = {
 335     {"__length_hint__", (PyCFunction)reversed_len, METH_NOARGS, length_hint_doc},
 336     {NULL,              NULL}           /* sentinel */
 337 };
 338
 339 PyTypeObject PyReversed_Type = {
 340     PyVarObject_HEAD_INIT(&PyType_Type, 0)
 341     "reversed",                     /* tp_name */
 342     sizeof(reversedobject),         /* tp_basicsize */
 343     0,                              /* tp_itemsize */
 344     /* methods */
 345     (destructor)reversed_dealloc,   /* tp_dealloc */
 346     0,                              /* tp_print */
 347     0,                              /* tp_getattr */
 348     0,                              /* tp_setattr */
 349     0,                              /* tp_compare */
 350     0,                              /* tp_repr */
 351     0,                              /* tp_as_number */
 352     0,                                  /* tp_as_sequence */
 353     0,                              /* tp_as_mapping */
 354     0,                              /* tp_hash */
 355     0,                              /* tp_call */
 356     0,                              /* tp_str */
 357     PyObject_GenericGetAttr,        /* tp_getattro */
 358     0,                              /* tp_setattro */
 359     0,                              /* tp_as_buffer */
 360     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 361         Py_TPFLAGS_BASETYPE,    /* tp_flags */
 362     reversed_doc,                   /* tp_doc */
 363     (traverseproc)reversed_traverse,/* tp_traverse */
 364     0,                              /* tp_clear */
 365     0,                              /* tp_richcompare */
 366     0,                              /* tp_weaklistoffset */
 367     PyObject_SelfIter,                  /* tp_iter */
 368     (iternextfunc)reversed_next,    /* tp_iternext */
 369     reversediter_methods,               /* tp_methods */
 370     0,                              /* tp_members */
 371     0,                              /* tp_getset */
 372     0,                              /* tp_base */
 373     0,                              /* tp_dict */
 374     0,                              /* tp_descr_get */
 375     0,                              /* tp_descr_set */
 376     0,                              /* tp_dictoffset */
 377     0,                              /* tp_init */
 378     PyType_GenericAlloc,            /* tp_alloc */
 379     reversed_new,                   /* tp_new */
 380     PyObject_GC_Del,                /* tp_free */
 381 };