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