Line data Source code
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 345206 : enum_new_impl(PyTypeObject *type, PyObject *iterable, PyObject *start)
45 : /*[clinic end generated code: output=e95e6e439f812c10 input=782e4911efcb8acf]*/
46 : {
47 : enumobject *en;
48 :
49 345206 : en = (enumobject *)type->tp_alloc(type, 0);
50 345206 : if (en == NULL)
51 0 : return NULL;
52 345206 : if (start != NULL) {
53 41707 : start = PyNumber_Index(start);
54 41707 : if (start == NULL) {
55 7 : Py_DECREF(en);
56 7 : return NULL;
57 : }
58 41700 : assert(PyLong_Check(start));
59 41700 : en->en_index = PyLong_AsSsize_t(start);
60 41700 : if (en->en_index == -1 && PyErr_Occurred()) {
61 33 : PyErr_Clear();
62 33 : en->en_index = PY_SSIZE_T_MAX;
63 33 : en->en_longindex = start;
64 : } else {
65 41667 : en->en_longindex = NULL;
66 41667 : Py_DECREF(start);
67 : }
68 : } else {
69 303499 : en->en_index = 0;
70 303499 : en->en_longindex = NULL;
71 : }
72 345199 : en->en_sit = PyObject_GetIter(iterable);
73 345199 : if (en->en_sit == NULL) {
74 26 : Py_DECREF(en);
75 26 : return NULL;
76 : }
77 345173 : en->en_result = PyTuple_Pack(2, Py_None, Py_None);
78 345173 : if (en->en_result == NULL) {
79 0 : Py_DECREF(en);
80 0 : return NULL;
81 : }
82 345173 : en->one = _PyLong_GetOne(); /* borrowed reference */
83 345173 : return (PyObject *)en;
84 : }
85 :
86 724 : static int check_keyword(PyObject *kwnames, int index,
87 : const char *name)
88 : {
89 724 : PyObject *kw = PyTuple_GET_ITEM(kwnames, index);
90 724 : if (!_PyUnicode_EqualToASCIIString(kw, name)) {
91 16 : PyErr_Format(PyExc_TypeError,
92 : "'%S' is an invalid keyword argument for enumerate()", kw);
93 16 : return 0;
94 : }
95 708 : return 1;
96 : }
97 :
98 : // TODO: Use AC when bpo-43447 is supported
99 : static PyObject *
100 345192 : enumerate_vectorcall(PyObject *type, PyObject *const *args,
101 : size_t nargsf, PyObject *kwnames)
102 : {
103 345192 : PyTypeObject *tp = _PyType_CAST(type);
104 345192 : Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
105 345192 : Py_ssize_t nkwargs = 0;
106 345192 : if (kwnames != NULL) {
107 716 : nkwargs = PyTuple_GET_SIZE(kwnames);
108 : }
109 :
110 : // Manually implement enumerate(iterable, start=...)
111 345192 : if (nargs + nkwargs == 2) {
112 41697 : if (nkwargs == 1) {
113 688 : if (!check_keyword(kwnames, 0, "start")) {
114 0 : return NULL;
115 : }
116 41009 : } else if (nkwargs == 2) {
117 20 : PyObject *kw0 = PyTuple_GET_ITEM(kwnames, 0);
118 20 : if (_PyUnicode_EqualToASCIIString(kw0, "start")) {
119 8 : if (!check_keyword(kwnames, 1, "iterable")) {
120 4 : return NULL;
121 : }
122 4 : return enum_new_impl(tp, args[1], args[0]);
123 : }
124 20 : if (!check_keyword(kwnames, 0, "iterable") ||
125 8 : !check_keyword(kwnames, 1, "start")) {
126 8 : return NULL;
127 : }
128 :
129 : }
130 41681 : return enum_new_impl(tp, args[0], args[1]);
131 : }
132 :
133 303495 : if (nargs + nkwargs == 1) {
134 303487 : if (nkwargs == 1 && !check_keyword(kwnames, 0, "iterable")) {
135 4 : return NULL;
136 : }
137 303483 : return enum_new_impl(tp, args[0], NULL);
138 : }
139 :
140 8 : if (nargs == 0) {
141 4 : PyErr_SetString(PyExc_TypeError,
142 : "enumerate() missing required argument 'iterable'");
143 4 : return NULL;
144 : }
145 :
146 4 : PyErr_Format(PyExc_TypeError,
147 : "enumerate() takes at most 2 arguments (%d given)", nargs + nkwargs);
148 4 : return NULL;
149 : }
150 :
151 : static void
152 345206 : enum_dealloc(enumobject *en)
153 : {
154 345206 : PyObject_GC_UnTrack(en);
155 345206 : Py_XDECREF(en->en_sit);
156 345206 : Py_XDECREF(en->en_result);
157 345206 : Py_XDECREF(en->en_longindex);
158 345206 : Py_TYPE(en)->tp_free(en);
159 345206 : }
160 :
161 : static int
162 70878 : enum_traverse(enumobject *en, visitproc visit, void *arg)
163 : {
164 70878 : Py_VISIT(en->en_sit);
165 70878 : Py_VISIT(en->en_result);
166 70878 : Py_VISIT(en->en_longindex);
167 70878 : return 0;
168 : }
169 :
170 : static PyObject *
171 52 : enum_next_long(enumobject *en, PyObject* next_item)
172 : {
173 52 : PyObject *result = en->en_result;
174 : PyObject *next_index;
175 : PyObject *stepped_up;
176 : PyObject *old_index;
177 : PyObject *old_item;
178 :
179 52 : if (en->en_longindex == NULL) {
180 0 : en->en_longindex = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
181 0 : if (en->en_longindex == NULL) {
182 0 : Py_DECREF(next_item);
183 0 : return NULL;
184 : }
185 : }
186 52 : next_index = en->en_longindex;
187 52 : assert(next_index != NULL);
188 52 : stepped_up = PyNumber_Add(next_index, en->one);
189 52 : if (stepped_up == NULL) {
190 0 : Py_DECREF(next_item);
191 0 : return NULL;
192 : }
193 52 : en->en_longindex = stepped_up;
194 :
195 52 : if (Py_REFCNT(result) == 1) {
196 24 : Py_INCREF(result);
197 24 : old_index = PyTuple_GET_ITEM(result, 0);
198 24 : old_item = PyTuple_GET_ITEM(result, 1);
199 24 : PyTuple_SET_ITEM(result, 0, next_index);
200 24 : PyTuple_SET_ITEM(result, 1, next_item);
201 24 : Py_DECREF(old_index);
202 24 : 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 24 : if (!_PyObject_GC_IS_TRACKED(result)) {
206 1 : _PyObject_GC_TRACK(result);
207 : }
208 24 : return result;
209 : }
210 28 : result = PyTuple_New(2);
211 28 : if (result == NULL) {
212 0 : Py_DECREF(next_index);
213 0 : Py_DECREF(next_item);
214 0 : return NULL;
215 : }
216 28 : PyTuple_SET_ITEM(result, 0, next_index);
217 28 : PyTuple_SET_ITEM(result, 1, next_item);
218 28 : return result;
219 : }
220 :
221 : static PyObject *
222 5168900 : enum_next(enumobject *en)
223 : {
224 : PyObject *next_index;
225 : PyObject *next_item;
226 5168900 : PyObject *result = en->en_result;
227 5168900 : PyObject *it = en->en_sit;
228 : PyObject *old_index;
229 : PyObject *old_item;
230 :
231 5168900 : next_item = (*Py_TYPE(it)->tp_iternext)(it);
232 5168900 : if (next_item == NULL)
233 342845 : return NULL;
234 :
235 4826060 : if (en->en_index == PY_SSIZE_T_MAX)
236 52 : return enum_next_long(en, next_item);
237 :
238 4826000 : next_index = PyLong_FromSsize_t(en->en_index);
239 4826000 : if (next_index == NULL) {
240 0 : Py_DECREF(next_item);
241 0 : return NULL;
242 : }
243 4826000 : en->en_index++;
244 :
245 4826000 : if (Py_REFCNT(result) == 1) {
246 4615850 : Py_INCREF(result);
247 4615850 : old_index = PyTuple_GET_ITEM(result, 0);
248 4615850 : old_item = PyTuple_GET_ITEM(result, 1);
249 4615850 : PyTuple_SET_ITEM(result, 0, next_index);
250 4615850 : PyTuple_SET_ITEM(result, 1, next_item);
251 4615850 : Py_DECREF(old_index);
252 4615850 : 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 4615850 : if (!_PyObject_GC_IS_TRACKED(result)) {
256 803 : _PyObject_GC_TRACK(result);
257 : }
258 4615850 : return result;
259 : }
260 210152 : result = PyTuple_New(2);
261 210152 : if (result == NULL) {
262 0 : Py_DECREF(next_index);
263 0 : Py_DECREF(next_item);
264 0 : return NULL;
265 : }
266 210152 : PyTuple_SET_ITEM(result, 0, next_index);
267 210152 : PyTuple_SET_ITEM(result, 1, next_item);
268 210152 : return result;
269 : }
270 :
271 : static PyObject *
272 78 : enum_reduce(enumobject *en, PyObject *Py_UNUSED(ignored))
273 : {
274 78 : if (en->en_longindex != NULL)
275 12 : return Py_BuildValue("O(OO)", Py_TYPE(en), en->en_sit, en->en_longindex);
276 : else
277 66 : 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 146381 : 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 146381 : reversed_meth = _PyObject_LookupSpecial(seq, &_Py_ID(__reversed__));
361 146381 : if (reversed_meth == Py_None) {
362 2 : Py_DECREF(reversed_meth);
363 2 : PyErr_Format(PyExc_TypeError,
364 : "'%.200s' object is not reversible",
365 2 : Py_TYPE(seq)->tp_name);
366 2 : return NULL;
367 : }
368 146379 : if (reversed_meth != NULL) {
369 137028 : PyObject *res = _PyObject_CallNoArgs(reversed_meth);
370 137028 : Py_DECREF(reversed_meth);
371 137028 : return res;
372 : }
373 9351 : else if (PyErr_Occurred())
374 1 : return NULL;
375 :
376 9350 : if (!PySequence_Check(seq)) {
377 11 : PyErr_Format(PyExc_TypeError,
378 : "'%.200s' object is not reversible",
379 11 : Py_TYPE(seq)->tp_name);
380 11 : return NULL;
381 : }
382 :
383 9339 : n = PySequence_Size(seq);
384 9339 : if (n == -1)
385 1 : return NULL;
386 :
387 9338 : ro = (reversedobject *)type->tp_alloc(type, 0);
388 9338 : if (ro == NULL)
389 0 : return NULL;
390 :
391 9338 : ro->index = n-1;
392 9338 : Py_INCREF(seq);
393 9338 : ro->seq = seq;
394 9338 : return (PyObject *)ro;
395 : }
396 :
397 : static PyObject *
398 146384 : reversed_vectorcall(PyObject *type, PyObject * const*args,
399 : size_t nargsf, PyObject *kwnames)
400 : {
401 146384 : if (!_PyArg_NoKwnames("reversed", kwnames)) {
402 1 : return NULL;
403 : }
404 :
405 146383 : Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
406 146383 : if (!_PyArg_CheckPositional("reversed", nargs, 1, 1)) {
407 2 : return NULL;
408 : }
409 :
410 146381 : return reversed_new_impl(_PyType_CAST(type), args[0]);
411 : }
412 :
413 : static void
414 9338 : reversed_dealloc(reversedobject *ro)
415 : {
416 9338 : PyObject_GC_UnTrack(ro);
417 9338 : Py_XDECREF(ro->seq);
418 9338 : Py_TYPE(ro)->tp_free(ro);
419 9338 : }
420 :
421 : static int
422 242 : reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
423 : {
424 242 : Py_VISIT(ro->seq);
425 242 : return 0;
426 : }
427 :
428 : static PyObject *
429 23239 : reversed_next(reversedobject *ro)
430 : {
431 : PyObject *item;
432 23239 : Py_ssize_t index = ro->index;
433 :
434 23239 : if (index >= 0) {
435 15588 : item = PySequence_GetItem(ro->seq, index);
436 15588 : if (item != NULL) {
437 15588 : ro->index--;
438 15588 : return item;
439 : }
440 0 : if (PyErr_ExceptionMatches(PyExc_IndexError) ||
441 0 : PyErr_ExceptionMatches(PyExc_StopIteration))
442 0 : PyErr_Clear();
443 : }
444 7651 : ro->index = -1;
445 7651 : Py_CLEAR(ro->seq);
446 7651 : return NULL;
447 : }
448 :
449 : static PyObject *
450 1662 : reversed_len(reversedobject *ro, PyObject *Py_UNUSED(ignored))
451 : {
452 : Py_ssize_t position, seqsize;
453 :
454 1662 : if (ro->seq == NULL)
455 15 : return PyLong_FromLong(0);
456 1647 : seqsize = PySequence_Size(ro->seq);
457 1647 : if (seqsize == -1)
458 1 : return NULL;
459 1646 : position = ro->index + 1;
460 1646 : 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 348 : reversed_reduce(reversedobject *ro, PyObject *Py_UNUSED(ignored))
467 : {
468 348 : if (ro->seq)
469 270 : return Py_BuildValue("O(O)n", Py_TYPE(ro), ro->seq, ro->index);
470 : else
471 78 : return Py_BuildValue("O(())", Py_TYPE(ro));
472 : }
473 :
474 : static PyObject *
475 282 : reversed_setstate(reversedobject *ro, PyObject *state)
476 : {
477 282 : Py_ssize_t index = PyLong_AsSsize_t(state);
478 282 : if (index == -1 && PyErr_Occurred())
479 0 : return NULL;
480 282 : if (ro->seq != 0) {
481 282 : Py_ssize_t n = PySequence_Size(ro->seq);
482 282 : if (n < 0)
483 0 : return NULL;
484 282 : if (index < -1)
485 0 : index = -1;
486 282 : else if (index > n-1)
487 0 : index = n-1;
488 282 : ro->index = index;
489 : }
490 282 : 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 : };
|
