1 /* enumerate object */
2
3 #include "Python.h"
4
5 #include "clinic/enumobject.c.h"
6
7 /*[clinic input]
8 class enumerate "enumobject *" "&PyEnum_Type"
9 class reversed "reversedobject *" "&PyReversed_Type"
10 [clinic start generated code]*/
11 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=d2dfdf1a88c88975]*/
12
13 typedef struct {
14 PyObject_HEAD
15 Py_ssize_t en_index; /* current index of enumeration */
16 PyObject* en_sit; /* secondary iterator of enumeration */
17 PyObject* en_result; /* result tuple */
18 PyObject* en_longindex; /* index for sequences >= PY_SSIZE_T_MAX */
19 } enumobject;
20
21
22 /*[clinic input]
23 @classmethod
24 enumerate.__new__ as enum_new
25
26 iterable: object
27 an object supporting iteration
28 start: object = 0
29
30 Return an enumerate object.
31
32 The enumerate object yields pairs containing a count (from start, which
33 defaults to zero) and a value yielded by the iterable argument.
34
35 enumerate is useful for obtaining an indexed list:
36 (0, seq[0]), (1, seq[1]), (2, seq[2]), ...
37 [clinic start generated code]*/
38
39 static PyObject *
enum_new_impl(PyTypeObject * type,PyObject * iterable,PyObject * start)40 enum_new_impl(PyTypeObject *type, PyObject *iterable, PyObject *start)
41 /*[clinic end generated code: output=e95e6e439f812c10 input=782e4911efcb8acf]*/
42 {
43 enumobject *en;
44
45 en = (enumobject *)type->tp_alloc(type, 0);
46 if (en == NULL)
47 return NULL;
48 if (start != NULL) {
49 start = PyNumber_Index(start);
50 if (start == NULL) {
51 Py_DECREF(en);
52 return NULL;
53 }
54 assert(PyLong_Check(start));
55 en->en_index = PyLong_AsSsize_t(start);
56 if (en->en_index == -1 && PyErr_Occurred()) {
57 PyErr_Clear();
58 en->en_index = PY_SSIZE_T_MAX;
59 en->en_longindex = start;
60 } else {
61 en->en_longindex = NULL;
62 Py_DECREF(start);
63 }
64 } else {
65 en->en_index = 0;
66 en->en_longindex = NULL;
67 }
68 en->en_sit = PyObject_GetIter(iterable);
69 if (en->en_sit == NULL) {
70 Py_DECREF(en);
71 return NULL;
72 }
73 en->en_result = PyTuple_Pack(2, Py_None, Py_None);
74 if (en->en_result == NULL) {
75 Py_DECREF(en);
76 return NULL;
77 }
78 return (PyObject *)en;
79 }
80
81 static void
enum_dealloc(enumobject * en)82 enum_dealloc(enumobject *en)
83 {
84 PyObject_GC_UnTrack(en);
85 Py_XDECREF(en->en_sit);
86 Py_XDECREF(en->en_result);
87 Py_XDECREF(en->en_longindex);
88 Py_TYPE(en)->tp_free(en);
89 }
90
91 static int
enum_traverse(enumobject * en,visitproc visit,void * arg)92 enum_traverse(enumobject *en, visitproc visit, void *arg)
93 {
94 Py_VISIT(en->en_sit);
95 Py_VISIT(en->en_result);
96 Py_VISIT(en->en_longindex);
97 return 0;
98 }
99
100 static PyObject *
enum_next_long(enumobject * en,PyObject * next_item)101 enum_next_long(enumobject *en, PyObject* next_item)
102 {
103 PyObject *result = en->en_result;
104 PyObject *next_index;
105 PyObject *stepped_up;
106 PyObject *old_index;
107 PyObject *old_item;
108
109 if (en->en_longindex == NULL) {
110 en->en_longindex = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
111 if (en->en_longindex == NULL) {
112 Py_DECREF(next_item);
113 return NULL;
114 }
115 }
116 next_index = en->en_longindex;
117 assert(next_index != NULL);
118 stepped_up = PyNumber_Add(next_index, _PyLong_One);
119 if (stepped_up == NULL) {
120 Py_DECREF(next_item);
121 return NULL;
122 }
123 en->en_longindex = stepped_up;
124
125 if (result->ob_refcnt == 1) {
126 Py_INCREF(result);
127 old_index = PyTuple_GET_ITEM(result, 0);
128 old_item = PyTuple_GET_ITEM(result, 1);
129 PyTuple_SET_ITEM(result, 0, next_index);
130 PyTuple_SET_ITEM(result, 1, next_item);
131 Py_DECREF(old_index);
132 Py_DECREF(old_item);
133 return result;
134 }
135 result = PyTuple_New(2);
136 if (result == NULL) {
137 Py_DECREF(next_index);
138 Py_DECREF(next_item);
139 return NULL;
140 }
141 PyTuple_SET_ITEM(result, 0, next_index);
142 PyTuple_SET_ITEM(result, 1, next_item);
143 return result;
144 }
145
146 static PyObject *
enum_next(enumobject * en)147 enum_next(enumobject *en)
148 {
149 PyObject *next_index;
150 PyObject *next_item;
151 PyObject *result = en->en_result;
152 PyObject *it = en->en_sit;
153 PyObject *old_index;
154 PyObject *old_item;
155
156 next_item = (*Py_TYPE(it)->tp_iternext)(it);
157 if (next_item == NULL)
158 return NULL;
159
160 if (en->en_index == PY_SSIZE_T_MAX)
161 return enum_next_long(en, next_item);
162
163 next_index = PyLong_FromSsize_t(en->en_index);
164 if (next_index == NULL) {
165 Py_DECREF(next_item);
166 return NULL;
167 }
168 en->en_index++;
169
170 if (result->ob_refcnt == 1) {
171 Py_INCREF(result);
172 old_index = PyTuple_GET_ITEM(result, 0);
173 old_item = PyTuple_GET_ITEM(result, 1);
174 PyTuple_SET_ITEM(result, 0, next_index);
175 PyTuple_SET_ITEM(result, 1, next_item);
176 Py_DECREF(old_index);
177 Py_DECREF(old_item);
178 return result;
179 }
180 result = PyTuple_New(2);
181 if (result == NULL) {
182 Py_DECREF(next_index);
183 Py_DECREF(next_item);
184 return NULL;
185 }
186 PyTuple_SET_ITEM(result, 0, next_index);
187 PyTuple_SET_ITEM(result, 1, next_item);
188 return result;
189 }
190
191 static PyObject *
enum_reduce(enumobject * en)192 enum_reduce(enumobject *en)
193 {
194 if (en->en_longindex != NULL)
195 return Py_BuildValue("O(OO)", Py_TYPE(en), en->en_sit, en->en_longindex);
196 else
197 return Py_BuildValue("O(On)", Py_TYPE(en), en->en_sit, en->en_index);
198 }
199
200 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
201
202 static PyMethodDef enum_methods[] = {
203 {"__reduce__", (PyCFunction)enum_reduce, METH_NOARGS, reduce_doc},
204 {NULL, NULL} /* sentinel */
205 };
206
207 PyTypeObject PyEnum_Type = {
208 PyVarObject_HEAD_INIT(&PyType_Type, 0)
209 "enumerate", /* tp_name */
210 sizeof(enumobject), /* tp_basicsize */
211 0, /* tp_itemsize */
212 /* methods */
213 (destructor)enum_dealloc, /* tp_dealloc */
214 0, /* tp_print */
215 0, /* tp_getattr */
216 0, /* tp_setattr */
217 0, /* tp_reserved */
218 0, /* tp_repr */
219 0, /* tp_as_number */
220 0, /* tp_as_sequence */
221 0, /* tp_as_mapping */
222 0, /* tp_hash */
223 0, /* tp_call */
224 0, /* tp_str */
225 PyObject_GenericGetAttr, /* tp_getattro */
226 0, /* tp_setattro */
227 0, /* tp_as_buffer */
228 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
229 Py_TPFLAGS_BASETYPE, /* tp_flags */
230 enum_new__doc__, /* tp_doc */
231 (traverseproc)enum_traverse, /* tp_traverse */
232 0, /* tp_clear */
233 0, /* tp_richcompare */
234 0, /* tp_weaklistoffset */
235 PyObject_SelfIter, /* tp_iter */
236 (iternextfunc)enum_next, /* tp_iternext */
237 enum_methods, /* tp_methods */
238 0, /* tp_members */
239 0, /* tp_getset */
240 0, /* tp_base */
241 0, /* tp_dict */
242 0, /* tp_descr_get */
243 0, /* tp_descr_set */
244 0, /* tp_dictoffset */
245 0, /* tp_init */
246 PyType_GenericAlloc, /* tp_alloc */
247 enum_new, /* tp_new */
248 PyObject_GC_Del, /* tp_free */
249 };
250
251 /* Reversed Object ***************************************************************/
252
253 typedef struct {
254 PyObject_HEAD
255 Py_ssize_t index;
256 PyObject* seq;
257 } reversedobject;
258
259 /*[clinic input]
260 @classmethod
261 reversed.__new__ as reversed_new
262
263 sequence as seq: object
264 /
265
266 Return a reverse iterator over the values of the given sequence.
267 [clinic start generated code]*/
268
269 static PyObject *
reversed_new_impl(PyTypeObject * type,PyObject * seq)270 reversed_new_impl(PyTypeObject *type, PyObject *seq)
271 /*[clinic end generated code: output=f7854cc1df26f570 input=aeb720361e5e3f1d]*/
272 {
273 Py_ssize_t n;
274 PyObject *reversed_meth;
275 reversedobject *ro;
276 _Py_IDENTIFIER(__reversed__);
277
278 reversed_meth = _PyObject_LookupSpecial(seq, &PyId___reversed__);
279 if (reversed_meth == Py_None) {
280 Py_DECREF(reversed_meth);
281 PyErr_Format(PyExc_TypeError,
282 "'%.200s' object is not reversible",
283 Py_TYPE(seq)->tp_name);
284 return NULL;
285 }
286 if (reversed_meth != NULL) {
287 PyObject *res = _PyObject_CallNoArg(reversed_meth);
288 Py_DECREF(reversed_meth);
289 return res;
290 }
291 else if (PyErr_Occurred())
292 return NULL;
293
294 if (!PySequence_Check(seq)) {
295 PyErr_Format(PyExc_TypeError,
296 "'%.200s' object is not reversible",
297 Py_TYPE(seq)->tp_name);
298 return NULL;
299 }
300
301 n = PySequence_Size(seq);
302 if (n == -1)
303 return NULL;
304
305 ro = (reversedobject *)type->tp_alloc(type, 0);
306 if (ro == NULL)
307 return NULL;
308
309 ro->index = n-1;
310 Py_INCREF(seq);
311 ro->seq = seq;
312 return (PyObject *)ro;
313 }
314
315 static void
reversed_dealloc(reversedobject * ro)316 reversed_dealloc(reversedobject *ro)
317 {
318 PyObject_GC_UnTrack(ro);
319 Py_XDECREF(ro->seq);
320 Py_TYPE(ro)->tp_free(ro);
321 }
322
323 static int
reversed_traverse(reversedobject * ro,visitproc visit,void * arg)324 reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
325 {
326 Py_VISIT(ro->seq);
327 return 0;
328 }
329
330 static PyObject *
reversed_next(reversedobject * ro)331 reversed_next(reversedobject *ro)
332 {
333 PyObject *item;
334 Py_ssize_t index = ro->index;
335
336 if (index >= 0) {
337 item = PySequence_GetItem(ro->seq, index);
338 if (item != NULL) {
339 ro->index--;
340 return item;
341 }
342 if (PyErr_ExceptionMatches(PyExc_IndexError) ||
343 PyErr_ExceptionMatches(PyExc_StopIteration))
344 PyErr_Clear();
345 }
346 ro->index = -1;
347 Py_CLEAR(ro->seq);
348 return NULL;
349 }
350
351 static PyObject *
reversed_len(reversedobject * ro)352 reversed_len(reversedobject *ro)
353 {
354 Py_ssize_t position, seqsize;
355
356 if (ro->seq == NULL)
357 return PyLong_FromLong(0);
358 seqsize = PySequence_Size(ro->seq);
359 if (seqsize == -1)
360 return NULL;
361 position = ro->index + 1;
362 return PyLong_FromSsize_t((seqsize < position) ? 0 : position);
363 }
364
365 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
366
367 static PyObject *
reversed_reduce(reversedobject * ro)368 reversed_reduce(reversedobject *ro)
369 {
370 if (ro->seq)
371 return Py_BuildValue("O(O)n", Py_TYPE(ro), ro->seq, ro->index);
372 else
373 return Py_BuildValue("O(())", Py_TYPE(ro));
374 }
375
376 static PyObject *
reversed_setstate(reversedobject * ro,PyObject * state)377 reversed_setstate(reversedobject *ro, PyObject *state)
378 {
379 Py_ssize_t index = PyLong_AsSsize_t(state);
380 if (index == -1 && PyErr_Occurred())
381 return NULL;
382 if (ro->seq != 0) {
383 Py_ssize_t n = PySequence_Size(ro->seq);
384 if (n < 0)
385 return NULL;
386 if (index < -1)
387 index = -1;
388 else if (index > n-1)
389 index = n-1;
390 ro->index = index;
391 }
392 Py_RETURN_NONE;
393 }
394
395 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
396
397 static PyMethodDef reversediter_methods[] = {
398 {"__length_hint__", (PyCFunction)reversed_len, METH_NOARGS, length_hint_doc},
399 {"__reduce__", (PyCFunction)reversed_reduce, METH_NOARGS, reduce_doc},
400 {"__setstate__", (PyCFunction)reversed_setstate, METH_O, setstate_doc},
401 {NULL, NULL} /* sentinel */
402 };
403
404 PyTypeObject PyReversed_Type = {
405 PyVarObject_HEAD_INIT(&PyType_Type, 0)
406 "reversed", /* tp_name */
407 sizeof(reversedobject), /* tp_basicsize */
408 0, /* tp_itemsize */
409 /* methods */
410 (destructor)reversed_dealloc, /* tp_dealloc */
411 0, /* tp_print */
412 0, /* tp_getattr */
413 0, /* tp_setattr */
414 0, /* tp_reserved */
415 0, /* tp_repr */
416 0, /* tp_as_number */
417 0, /* tp_as_sequence */
418 0, /* tp_as_mapping */
419 0, /* tp_hash */
420 0, /* tp_call */
421 0, /* tp_str */
422 PyObject_GenericGetAttr, /* tp_getattro */
423 0, /* tp_setattro */
424 0, /* tp_as_buffer */
425 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
426 Py_TPFLAGS_BASETYPE, /* tp_flags */
427 reversed_new__doc__, /* tp_doc */
428 (traverseproc)reversed_traverse,/* tp_traverse */
429 0, /* tp_clear */
430 0, /* tp_richcompare */
431 0, /* tp_weaklistoffset */
432 PyObject_SelfIter, /* tp_iter */
433 (iternextfunc)reversed_next, /* tp_iternext */
434 reversediter_methods, /* tp_methods */
435 0, /* tp_members */
436 0, /* tp_getset */
437 0, /* tp_base */
438 0, /* tp_dict */
439 0, /* tp_descr_get */
440 0, /* tp_descr_set */
441 0, /* tp_dictoffset */
442 0, /* tp_init */
443 PyType_GenericAlloc, /* tp_alloc */
444 reversed_new, /* tp_new */
445 PyObject_GC_Del, /* tp_free */
446 };
447