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]
10 class enumerate "enumobject *" "&PyEnum_Type"
11 class reversed "reversedobject *" "&PyReversed_Type"
12 [clinic start generated code]*/
13 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=d2dfdf1a88c88975]*/
14
15 typedef 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
26 enumerate.__new__ as enum_new
27
28 iterable: object
29 an object supporting iteration
30 start: object = 0
31
32 Return an enumerate object.
33
34 The enumerate object yields pairs containing a count (from start, which
35 defaults to zero) and a value yielded by the iterable argument.
36
37 enumerate is useful for obtaining an indexed list:
38 (0, seq[0]), (1, seq[1]), (2, seq[2]), ...
39 [clinic start generated code]*/
40
41 static PyObject *
enum_new_impl(PyTypeObject * type,PyObject * iterable,PyObject * start)42 enum_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
83 static void
enum_dealloc(enumobject * en)84 enum_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
93 static int
enum_traverse(enumobject * en,visitproc visit,void * arg)94 enum_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
102 static PyObject *
enum_next_long(enumobject * en,PyObject * next_item)103 enum_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
153 static PyObject *
enum_next(enumobject * en)154 enum_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
203 static PyObject *
enum_reduce(enumobject * en,PyObject * Py_UNUSED (ignored))204 enum_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
212 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
213
214 static 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
221 PyTypeObject 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
267 typedef struct {
268 PyObject_HEAD
269 Py_ssize_t index;
270 PyObject* seq;
271 } reversedobject;
272
273 /*[clinic input]
274 @classmethod
275 reversed.__new__ as reversed_new
276
277 sequence as seq: object
278 /
279
280 Return a reverse iterator over the values of the given sequence.
281 [clinic start generated code]*/
282
283 static PyObject *
reversed_new_impl(PyTypeObject * type,PyObject * seq)284 reversed_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
329 static PyObject *
reversed_vectorcall(PyObject * type,PyObject * const * args,size_t nargsf,PyObject * kwnames)330 reversed_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
347 static void
reversed_dealloc(reversedobject * ro)348 reversed_dealloc(reversedobject *ro)
349 {
350 PyObject_GC_UnTrack(ro);
351 Py_XDECREF(ro->seq);
352 Py_TYPE(ro)->tp_free(ro);
353 }
354
355 static int
reversed_traverse(reversedobject * ro,visitproc visit,void * arg)356 reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
357 {
358 Py_VISIT(ro->seq);
359 return 0;
360 }
361
362 static PyObject *
reversed_next(reversedobject * ro)363 reversed_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
383 static PyObject *
reversed_len(reversedobject * ro,PyObject * Py_UNUSED (ignored))384 reversed_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
397 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
398
399 static PyObject *
reversed_reduce(reversedobject * ro,PyObject * Py_UNUSED (ignored))400 reversed_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
408 static PyObject *
reversed_setstate(reversedobject * ro,PyObject * state)409 reversed_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
427 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
428
429 static 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
436 PyTypeObject 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