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