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