Home | History | Annotate | Download | only in Objects
      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[] = {"iterable", "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(PyLong_Check(start));
     35         en->en_index = PyLong_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 = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
     90         if (en->en_longindex == NULL)
     91             return NULL;
     92     }
     93     if (one == NULL) {
     94         one = PyLong_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 = PyLong_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 static PyObject *
    162 enum_reduce(enumobject *en)
    163 {
    164     if (en->en_longindex != NULL)
    165         return Py_BuildValue("O(OO)", Py_TYPE(en), en->en_sit, en->en_longindex);
    166     else
    167         return Py_BuildValue("O(On)", Py_TYPE(en), en->en_sit, en->en_index);
    168 }
    169 
    170 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
    171 
    172 static PyMethodDef enum_methods[] = {
    173     {"__reduce__", (PyCFunction)enum_reduce, METH_NOARGS, reduce_doc},
    174     {NULL,              NULL}           /* sentinel */
    175 };
    176 
    177 PyDoc_STRVAR(enum_doc,
    178 "enumerate(iterable[, start]) -> iterator for index, value of iterable\n"
    179 "\n"
    180 "Return an enumerate object.  iterable must be another object that supports\n"
    181 "iteration.  The enumerate object yields pairs containing a count (from\n"
    182 "start, which defaults to zero) and a value yielded by the iterable argument.\n"
    183 "enumerate is useful for obtaining an indexed list:\n"
    184 "    (0, seq[0]), (1, seq[1]), (2, seq[2]), ...");
    185 
    186 PyTypeObject PyEnum_Type = {
    187     PyVarObject_HEAD_INIT(&PyType_Type, 0)
    188     "enumerate",                    /* tp_name */
    189     sizeof(enumobject),             /* tp_basicsize */
    190     0,                              /* tp_itemsize */
    191     /* methods */
    192     (destructor)enum_dealloc,       /* tp_dealloc */
    193     0,                              /* tp_print */
    194     0,                              /* tp_getattr */
    195     0,                              /* tp_setattr */
    196     0,                              /* tp_reserved */
    197     0,                              /* tp_repr */
    198     0,                              /* tp_as_number */
    199     0,                              /* tp_as_sequence */
    200     0,                              /* tp_as_mapping */
    201     0,                              /* tp_hash */
    202     0,                              /* tp_call */
    203     0,                              /* tp_str */
    204     PyObject_GenericGetAttr,        /* tp_getattro */
    205     0,                              /* tp_setattro */
    206     0,                              /* tp_as_buffer */
    207     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
    208         Py_TPFLAGS_BASETYPE,    /* tp_flags */
    209     enum_doc,                       /* tp_doc */
    210     (traverseproc)enum_traverse,    /* tp_traverse */
    211     0,                              /* tp_clear */
    212     0,                              /* tp_richcompare */
    213     0,                              /* tp_weaklistoffset */
    214     PyObject_SelfIter,                  /* tp_iter */
    215     (iternextfunc)enum_next,        /* tp_iternext */
    216     enum_methods,                   /* tp_methods */
    217     0,                              /* tp_members */
    218     0,                              /* tp_getset */
    219     0,                              /* tp_base */
    220     0,                              /* tp_dict */
    221     0,                              /* tp_descr_get */
    222     0,                              /* tp_descr_set */
    223     0,                              /* tp_dictoffset */
    224     0,                              /* tp_init */
    225     PyType_GenericAlloc,            /* tp_alloc */
    226     enum_new,                       /* tp_new */
    227     PyObject_GC_Del,                /* tp_free */
    228 };
    229 
    230 /* Reversed Object ***************************************************************/
    231 
    232 typedef struct {
    233     PyObject_HEAD
    234     Py_ssize_t      index;
    235     PyObject* seq;
    236 } reversedobject;
    237 
    238 static PyObject *
    239 reversed_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
    240 {
    241     Py_ssize_t n;
    242     PyObject *seq, *reversed_meth;
    243     reversedobject *ro;
    244     _Py_IDENTIFIER(__reversed__);
    245 
    246     if (type == &PyReversed_Type && !_PyArg_NoKeywords("reversed()", kwds))
    247         return NULL;
    248 
    249     if (!PyArg_UnpackTuple(args, "reversed", 1, 1, &seq) )
    250         return NULL;
    251 
    252     reversed_meth = _PyObject_LookupSpecial(seq, &PyId___reversed__);
    253     if (reversed_meth == Py_None) {
    254         Py_DECREF(reversed_meth);
    255         PyErr_Format(PyExc_TypeError,
    256                      "'%.200s' object is not reversible",
    257                      Py_TYPE(seq)->tp_name);
    258         return NULL;
    259     }
    260     if (reversed_meth != NULL) {
    261         PyObject *res = PyObject_CallFunctionObjArgs(reversed_meth, NULL);
    262         Py_DECREF(reversed_meth);
    263         return res;
    264     }
    265     else if (PyErr_Occurred())
    266         return NULL;
    267 
    268     if (!PySequence_Check(seq)) {
    269         PyErr_Format(PyExc_TypeError,
    270                      "'%.200s' object is not reversible",
    271                      Py_TYPE(seq)->tp_name);
    272         return NULL;
    273     }
    274 
    275     n = PySequence_Size(seq);
    276     if (n == -1)
    277         return NULL;
    278 
    279     ro = (reversedobject *)type->tp_alloc(type, 0);
    280     if (ro == NULL)
    281         return NULL;
    282 
    283     ro->index = n-1;
    284     Py_INCREF(seq);
    285     ro->seq = seq;
    286     return (PyObject *)ro;
    287 }
    288 
    289 static void
    290 reversed_dealloc(reversedobject *ro)
    291 {
    292     PyObject_GC_UnTrack(ro);
    293     Py_XDECREF(ro->seq);
    294     Py_TYPE(ro)->tp_free(ro);
    295 }
    296 
    297 static int
    298 reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
    299 {
    300     Py_VISIT(ro->seq);
    301     return 0;
    302 }
    303 
    304 static PyObject *
    305 reversed_next(reversedobject *ro)
    306 {
    307     PyObject *item;
    308     Py_ssize_t index = ro->index;
    309 
    310     if (index >= 0) {
    311         item = PySequence_GetItem(ro->seq, index);
    312         if (item != NULL) {
    313             ro->index--;
    314             return item;
    315         }
    316         if (PyErr_ExceptionMatches(PyExc_IndexError) ||
    317             PyErr_ExceptionMatches(PyExc_StopIteration))
    318             PyErr_Clear();
    319     }
    320     ro->index = -1;
    321     Py_CLEAR(ro->seq);
    322     return NULL;
    323 }
    324 
    325 PyDoc_STRVAR(reversed_doc,
    326 "reversed(sequence) -> reverse iterator over values of the sequence\n"
    327 "\n"
    328 "Return a reverse iterator");
    329 
    330 static PyObject *
    331 reversed_len(reversedobject *ro)
    332 {
    333     Py_ssize_t position, seqsize;
    334 
    335     if (ro->seq == NULL)
    336         return PyLong_FromLong(0);
    337     seqsize = PySequence_Size(ro->seq);
    338     if (seqsize == -1)
    339         return NULL;
    340     position = ro->index + 1;
    341     return PyLong_FromSsize_t((seqsize < position)  ?  0  :  position);
    342 }
    343 
    344 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
    345 
    346 static PyObject *
    347 reversed_reduce(reversedobject *ro)
    348 {
    349     if (ro->seq)
    350         return Py_BuildValue("O(O)n", Py_TYPE(ro), ro->seq, ro->index);
    351     else
    352         return Py_BuildValue("O(())", Py_TYPE(ro));
    353 }
    354 
    355 static PyObject *
    356 reversed_setstate(reversedobject *ro, PyObject *state)
    357 {
    358     Py_ssize_t index = PyLong_AsSsize_t(state);
    359     if (index == -1 && PyErr_Occurred())
    360         return NULL;
    361     if (ro->seq != 0) {
    362         Py_ssize_t n = PySequence_Size(ro->seq);
    363         if (n < 0)
    364             return NULL;
    365         if (index < -1)
    366             index = -1;
    367         else if (index > n-1)
    368             index = n-1;
    369         ro->index = index;
    370     }
    371     Py_RETURN_NONE;
    372 }
    373 
    374 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
    375 
    376 static PyMethodDef reversediter_methods[] = {
    377     {"__length_hint__", (PyCFunction)reversed_len, METH_NOARGS, length_hint_doc},
    378     {"__reduce__", (PyCFunction)reversed_reduce, METH_NOARGS, reduce_doc},
    379     {"__setstate__", (PyCFunction)reversed_setstate, METH_O, setstate_doc},
    380     {NULL,              NULL}           /* sentinel */
    381 };
    382 
    383 PyTypeObject PyReversed_Type = {
    384     PyVarObject_HEAD_INIT(&PyType_Type, 0)
    385     "reversed",                     /* tp_name */
    386     sizeof(reversedobject),         /* tp_basicsize */
    387     0,                              /* tp_itemsize */
    388     /* methods */
    389     (destructor)reversed_dealloc,   /* tp_dealloc */
    390     0,                              /* tp_print */
    391     0,                              /* tp_getattr */
    392     0,                              /* tp_setattr */
    393     0,                              /* tp_reserved */
    394     0,                              /* tp_repr */
    395     0,                              /* tp_as_number */
    396     0,                                  /* tp_as_sequence */
    397     0,                              /* tp_as_mapping */
    398     0,                              /* tp_hash */
    399     0,                              /* tp_call */
    400     0,                              /* tp_str */
    401     PyObject_GenericGetAttr,        /* tp_getattro */
    402     0,                              /* tp_setattro */
    403     0,                              /* tp_as_buffer */
    404     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
    405         Py_TPFLAGS_BASETYPE,    /* tp_flags */
    406     reversed_doc,                   /* tp_doc */
    407     (traverseproc)reversed_traverse,/* tp_traverse */
    408     0,                              /* tp_clear */
    409     0,                              /* tp_richcompare */
    410     0,                              /* tp_weaklistoffset */
    411     PyObject_SelfIter,                  /* tp_iter */
    412     (iternextfunc)reversed_next,    /* tp_iternext */
    413     reversediter_methods,               /* tp_methods */
    414     0,                              /* tp_members */
    415     0,                              /* tp_getset */
    416     0,                              /* tp_base */
    417     0,                              /* tp_dict */
    418     0,                              /* tp_descr_get */
    419     0,                              /* tp_descr_set */
    420     0,                              /* tp_dictoffset */
    421     0,                              /* tp_init */
    422     PyType_GenericAlloc,            /* tp_alloc */
    423     reversed_new,                   /* tp_new */
    424     PyObject_GC_Del,                /* tp_free */
    425 };
    426