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