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