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[] = {"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 PyDoc_STRVAR(enum_doc,
 162 "enumerate(iterable) -> 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 "zero) and a value yielded by the iterable argument.  enumerate is useful\n"
 167 "for obtaining an indexed list: (0, seq[0]), (1, seq[1]), (2, seq[2]), ...");
 168
 169 PyTypeObject PyEnum_Type = {
 170         PyVarObject_HEAD_INIT(&PyType_Type, 0)
 171         "enumerate",                    /* tp_name */
 172         sizeof(enumobject),             /* tp_basicsize */
 173         0,                              /* tp_itemsize */
 174         /* methods */
 175         (destructor)enum_dealloc,       /* tp_dealloc */
 176         0,                              /* tp_print */
 177         0,                              /* tp_getattr */
 178         0,                              /* tp_setattr */
 179         0,                              /* tp_reserved */
 180         0,                              /* tp_repr */
 181         0,                              /* tp_as_number */
 182         0,                              /* tp_as_sequence */
 183         0,                              /* tp_as_mapping */
 184         0,                              /* tp_hash */
 185         0,                              /* tp_call */
 186         0,                              /* tp_str */
 187         PyObject_GenericGetAttr,        /* tp_getattro */
 188         0,                              /* tp_setattro */
 189         0,                              /* tp_as_buffer */
 190         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 191                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 192         enum_doc,                       /* tp_doc */
 193         (traverseproc)enum_traverse,    /* tp_traverse */
 194         0,                              /* tp_clear */
 195         0,                              /* tp_richcompare */
 196         0,                              /* tp_weaklistoffset */
 197         PyObject_SelfIter,              /* tp_iter */
 198         (iternextfunc)enum_next,        /* tp_iternext */
 199         0,                              /* tp_methods */
 200         0,                              /* tp_members */
 201         0,                              /* tp_getset */
 202         0,                              /* tp_base */
 203         0,                              /* tp_dict */
 204         0,                              /* tp_descr_get */
 205         0,                              /* tp_descr_set */
 206         0,                              /* tp_dictoffset */
 207         0,                              /* tp_init */
 208         PyType_GenericAlloc,            /* tp_alloc */
 209         enum_new,                       /* tp_new */
 210         PyObject_GC_Del,                /* tp_free */
 211 };
 212
 213 /* Reversed Object ***************************************************************/
 214
 215 typedef struct {
 216         PyObject_HEAD
 217         Py_ssize_t      index;
 218         PyObject* seq;
 219 } reversedobject;
 220
 221 static PyObject *
 222 reversed_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 223 {
 224         Py_ssize_t n;
 225         PyObject *seq, *reversed_meth;
 226         static PyObject *reversed_cache = NULL;
 227         reversedobject *ro;
 228
 229         if (type == &PyReversed_Type && !_PyArg_NoKeywords("reversed()", kwds))
 230                 return NULL;
 231
 232         if (!PyArg_UnpackTuple(args, "reversed", 1, 1, &seq) )
 233                 return NULL;
 234
 235         reversed_meth = _PyObject_LookupSpecial(seq, "__reversed__", &reversed_cache);
 236         if (reversed_meth != NULL) {
 237                 PyObject *res = PyObject_CallFunctionObjArgs(reversed_meth, NULL);
 238                 Py_DECREF(reversed_meth);
 239                 return res;
 240         }
 241         else if (PyErr_Occurred())
 242                 return NULL;
 243
 244         if (!PySequence_Check(seq)) {
 245                 PyErr_SetString(PyExc_TypeError,
 246                                 "argument to reversed() must be a sequence");
 247                 return NULL;
 248         }
 249
 250         n = PySequence_Size(seq);
 251         if (n == -1)
 252                 return NULL;
 253
 254         ro = (reversedobject *)type->tp_alloc(type, 0);
 255         if (ro == NULL)
 256                 return NULL;
 257
 258         ro->index = n-1;
 259         Py_INCREF(seq);
 260         ro->seq = seq;
 261         return (PyObject *)ro;
 262 }
 263
 264 static void
 265 reversed_dealloc(reversedobject *ro)
 266 {
 267         PyObject_GC_UnTrack(ro);
 268         Py_XDECREF(ro->seq);
 269         Py_TYPE(ro)->tp_free(ro);
 270 }
 271
 272 static int
 273 reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
 274 {
 275         Py_VISIT(ro->seq);
 276         return 0;
 277 }
 278
 279 static PyObject *
 280 reversed_next(reversedobject *ro)
 281 {
 282         PyObject *item;
 283         Py_ssize_t index = ro->index;
 284
 285         if (index >= 0) {
 286                 item = PySequence_GetItem(ro->seq, index);
 287                 if (item != NULL) {
 288                         ro->index--;
 289                         return item;
 290                 }
 291                 if (PyErr_ExceptionMatches(PyExc_IndexError) ||
 292                     PyErr_ExceptionMatches(PyExc_StopIteration))
 293                         PyErr_Clear();
 294         }
 295         ro->index = -1;
 296         Py_CLEAR(ro->seq);
 297         return NULL;
 298 }
 299
 300 PyDoc_STRVAR(reversed_doc,
 301 "reversed(sequence) -> reverse iterator over values of the sequence\n"
 302 "\n"
 303 "Return a reverse iterator");
 304
 305 static PyObject *
 306 reversed_len(reversedobject *ro)
 307 {
 308         Py_ssize_t position, seqsize;
 309
 310         if (ro->seq == NULL)
 311                 return PyLong_FromLong(0);
 312         seqsize = PySequence_Size(ro->seq);
 313         if (seqsize == -1)
 314                 return NULL;
 315         position = ro->index + 1;
 316         return PyLong_FromSsize_t((seqsize < position)  ?  0  :  position);
 317 }
 318
 319 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
 320
 321 static PyMethodDef reversediter_methods[] = {
 322         {"__length_hint__", (PyCFunction)reversed_len, METH_NOARGS, length_hint_doc},
 323         {NULL,          NULL}           /* sentinel */
 324 };
 325
 326 PyTypeObject PyReversed_Type = {
 327         PyVarObject_HEAD_INIT(&PyType_Type, 0)
 328         "reversed",                     /* tp_name */
 329         sizeof(reversedobject),         /* tp_basicsize */
 330         0,                              /* tp_itemsize */
 331         /* methods */
 332         (destructor)reversed_dealloc,   /* tp_dealloc */
 333         0,                              /* tp_print */
 334         0,                              /* tp_getattr */
 335         0,                              /* tp_setattr */
 336         0,                              /* tp_reserved */
 337         0,                              /* tp_repr */
 338         0,                              /* tp_as_number */
 339         0,                              /* tp_as_sequence */
 340         0,                              /* tp_as_mapping */
 341         0,                              /* tp_hash */
 342         0,                              /* tp_call */
 343         0,                              /* tp_str */
 344         PyObject_GenericGetAttr,        /* tp_getattro */
 345         0,                              /* tp_setattro */
 346         0,                              /* tp_as_buffer */
 347         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 348                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 349         reversed_doc,                   /* tp_doc */
 350         (traverseproc)reversed_traverse,/* tp_traverse */
 351         0,                              /* tp_clear */
 352         0,                              /* tp_richcompare */
 353         0,                              /* tp_weaklistoffset */
 354         PyObject_SelfIter,              /* tp_iter */
 355         (iternextfunc)reversed_next,    /* tp_iternext */
 356         reversediter_methods,           /* tp_methods */
 357         0,                              /* tp_members */
 358         0,                              /* tp_getset */
 359         0,                              /* tp_base */
 360         0,                              /* tp_dict */
 361         0,                              /* tp_descr_get */
 362         0,                              /* tp_descr_set */
 363         0,                              /* tp_dictoffset */
 364         0,                              /* tp_init */
 365         PyType_GenericAlloc,            /* tp_alloc */
 366         reversed_new,                   /* tp_new */
 367         PyObject_GC_Del,                /* tp_free */
 368 };