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