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