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) -> iterator for index, value of iterable\n"
 163 "\n"
 164 "Return an enumerate object.  iterable must be an other 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_compare */
 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;
 226         reversedobject *ro;
 227
 228         if (type == &PyReversed_Type && !_PyArg_NoKeywords("reversed()", kwds))
 229                 return NULL;
 230
 231         if (!PyArg_UnpackTuple(args, "reversed", 1, 1, &seq) )
 232                 return NULL;
 233
 234         if (PyObject_HasAttrString(seq, "__reversed__"))
 235                 return PyObject_CallMethod(seq, "__reversed__", NULL);
 236
 237         if (!PySequence_Check(seq)) {
 238                 PyErr_SetString(PyExc_TypeError,
 239                                 "argument to reversed() must be a sequence");
 240                 return NULL;
 241         }
 242
 243         n = PySequence_Size(seq);
 244         if (n == -1)
 245                 return NULL;
 246
 247         ro = (reversedobject *)type->tp_alloc(type, 0);
 248         if (ro == NULL)
 249                 return NULL;
 250
 251         ro->index = n-1;
 252         Py_INCREF(seq);
 253         ro->seq = seq;
 254         return (PyObject *)ro;
 255 }
 256
 257 static void
 258 reversed_dealloc(reversedobject *ro)
 259 {
 260         PyObject_GC_UnTrack(ro);
 261         Py_XDECREF(ro->seq);
 262         Py_TYPE(ro)->tp_free(ro);
 263 }
 264
 265 static int
 266 reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
 267 {
 268         Py_VISIT(ro->seq);
 269         return 0;
 270 }
 271
 272 static PyObject *
 273 reversed_next(reversedobject *ro)
 274 {
 275         PyObject *item;
 276         Py_ssize_t index = ro->index;
 277
 278         if (index >= 0) {
 279                 item = PySequence_GetItem(ro->seq, index);
 280                 if (item != NULL) {
 281                         ro->index--;
 282                         return item;
 283                 }
 284                 if (PyErr_ExceptionMatches(PyExc_IndexError) ||
 285                     PyErr_ExceptionMatches(PyExc_StopIteration))
 286                         PyErr_Clear();
 287         }
 288         ro->index = -1;
 289         Py_CLEAR(ro->seq);
 290         return NULL;
 291 }
 292
 293 PyDoc_STRVAR(reversed_doc,
 294 "reversed(sequence) -> reverse iterator over values of the sequence\n"
 295 "\n"
 296 "Return a reverse iterator");
 297
 298 static PyObject *
 299 reversed_len(reversedobject *ro)
 300 {
 301         Py_ssize_t position, seqsize;
 302
 303         if (ro->seq == NULL)
 304                 return PyInt_FromLong(0);
 305         seqsize = PySequence_Size(ro->seq);
 306         if (seqsize == -1)
 307                 return NULL;
 308         position = ro->index + 1;
 309         return PyInt_FromSsize_t((seqsize < position)  ?  0  :  position);
 310 }
 311
 312 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
 313
 314 static PyMethodDef reversediter_methods[] = {
 315         {"__length_hint__", (PyCFunction)reversed_len, METH_NOARGS, length_hint_doc},
 316         {NULL,          NULL}           /* sentinel */
 317 };
 318
 319 PyTypeObject PyReversed_Type = {
 320         PyVarObject_HEAD_INIT(&PyType_Type, 0)
 321         "reversed",                     /* tp_name */
 322         sizeof(reversedobject),         /* tp_basicsize */
 323         0,                              /* tp_itemsize */
 324         /* methods */
 325         (destructor)reversed_dealloc,   /* tp_dealloc */
 326         0,                              /* tp_print */
 327         0,                              /* tp_getattr */
 328         0,                              /* tp_setattr */
 329         0,                              /* tp_compare */
 330         0,                              /* tp_repr */
 331         0,                              /* tp_as_number */
 332         0,                              /* tp_as_sequence */
 333         0,                              /* tp_as_mapping */
 334         0,                              /* tp_hash */
 335         0,                              /* tp_call */
 336         0,                              /* tp_str */
 337         PyObject_GenericGetAttr,        /* tp_getattro */
 338         0,                              /* tp_setattro */
 339         0,                              /* tp_as_buffer */
 340         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 341                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 342         reversed_doc,                   /* tp_doc */
 343         (traverseproc)reversed_traverse,/* tp_traverse */
 344         0,                              /* tp_clear */
 345         0,                              /* tp_richcompare */
 346         0,                              /* tp_weaklistoffset */
 347         PyObject_SelfIter,              /* tp_iter */
 348         (iternextfunc)reversed_next,    /* tp_iternext */
 349         reversediter_methods,           /* tp_methods */
 350         0,                              /* tp_members */
 351         0,                              /* tp_getset */
 352         0,                              /* tp_base */
 353         0,                              /* tp_dict */
 354         0,                              /* tp_descr_get */
 355         0,                              /* tp_descr_set */
 356         0,                              /* tp_dictoffset */
 357         0,                              /* tp_init */
 358         PyType_GenericAlloc,            /* tp_alloc */
 359         reversed_new,                   /* tp_new */
 360         PyObject_GC_Del,                /* tp_free */
 361 };