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[] = {"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 };
    382