1 /* enumerate object */
2
3 #include "Python.h"
4
5 typedef struct {
6 PyObject_HEAD
7 Py_ssize_t en_index; /* current index of enumeration */
8 PyObject* en_sit; /* secondary iterator of enumeration */
9 PyObject* en_result; /* result tuple */
10 PyObject* en_longindex; /* index for sequences >= PY_SSIZE_T_MAX */
11 } enumobject;
12
13 static PyObject *
enum_new(PyTypeObject * type,PyObject * args,PyObject * kwds)14 enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
15 {
16 enumobject *en;
17 PyObject *seq = NULL;
18 PyObject *start = NULL;
19 static char *kwlist[] = {"sequence", "start", 0};
20
21 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:enumerate", kwlist,
22 &seq, &start))
23 return NULL;
24
25 en = (enumobject *)type->tp_alloc(type, 0);
26 if (en == NULL)
27 return NULL;
28 if (start != NULL) {
29 start = PyNumber_Index(start);
30 if (start == NULL) {
31 Py_DECREF(en);
32 return NULL;
33 }
34 assert(_PyAnyInt_Check(start));
35 en->en_index = PyInt_AsSsize_t(start);
36 if (en->en_index == -1 && PyErr_Occurred()) {
37 PyErr_Clear();
38 en->en_index = PY_SSIZE_T_MAX;
39 en->en_longindex = start;
40 } else {
41 en->en_longindex = NULL;
42 Py_DECREF(start);
43 }
44 } else {
45 en->en_index = 0;
46 en->en_longindex = NULL;
47 }
48 en->en_sit = PyObject_GetIter(seq);
49 if (en->en_sit == NULL) {
50 Py_DECREF(en);
51 return NULL;
52 }
53 en->en_result = PyTuple_Pack(2, Py_None, Py_None);
54 if (en->en_result == NULL) {
55 Py_DECREF(en);
56 return NULL;
57 }
58 return (PyObject *)en;
59 }
60
61 static void
enum_dealloc(enumobject * en)62 enum_dealloc(enumobject *en)
63 {
64 PyObject_GC_UnTrack(en);
65 Py_XDECREF(en->en_sit);
66 Py_XDECREF(en->en_result);
67 Py_XDECREF(en->en_longindex);
68 Py_TYPE(en)->tp_free(en);
69 }
70
71 static int
enum_traverse(enumobject * en,visitproc visit,void * arg)72 enum_traverse(enumobject *en, visitproc visit, void *arg)
73 {
74 Py_VISIT(en->en_sit);
75 Py_VISIT(en->en_result);
76 Py_VISIT(en->en_longindex);
77 return 0;
78 }
79
80 static PyObject *
enum_next_long(enumobject * en,PyObject * next_item)81 enum_next_long(enumobject *en, PyObject* next_item)
82 {
83 static PyObject *one = NULL;
84 PyObject *result = en->en_result;
85 PyObject *next_index;
86 PyObject *stepped_up;
87
88 if (en->en_longindex == NULL) {
89 en->en_longindex = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
90 if (en->en_longindex == NULL) {
91 Py_DECREF(next_item);
92 return NULL;
93 }
94 }
95 if (one == NULL) {
96 one = PyInt_FromLong(1);
97 if (one == NULL) {
98 Py_DECREF(next_item);
99 return NULL;
100 }
101 }
102 next_index = en->en_longindex;
103 assert(next_index != NULL);
104 stepped_up = PyNumber_Add(next_index, one);
105 if (stepped_up == NULL) {
106 Py_DECREF(next_item);
107 return NULL;
108 }
109 en->en_longindex = stepped_up;
110
111 if (result->ob_refcnt == 1) {
112 Py_INCREF(result);
113 Py_DECREF(PyTuple_GET_ITEM(result, 0));
114 Py_DECREF(PyTuple_GET_ITEM(result, 1));
115 } else {
116 result = PyTuple_New(2);
117 if (result == NULL) {
118 Py_DECREF(next_index);
119 Py_DECREF(next_item);
120 return NULL;
121 }
122 }
123 PyTuple_SET_ITEM(result, 0, next_index);
124 PyTuple_SET_ITEM(result, 1, next_item);
125 return result;
126 }
127
128 static PyObject *
enum_next(enumobject * en)129 enum_next(enumobject *en)
130 {
131 PyObject *next_index;
132 PyObject *next_item;
133 PyObject *result = en->en_result;
134 PyObject *it = en->en_sit;
135
136 next_item = (*Py_TYPE(it)->tp_iternext)(it);
137 if (next_item == NULL)
138 return NULL;
139
140 if (en->en_index == PY_SSIZE_T_MAX)
141 return enum_next_long(en, next_item);
142
143 next_index = PyInt_FromSsize_t(en->en_index);
144 if (next_index == NULL) {
145 Py_DECREF(next_item);
146 return NULL;
147 }
148 en->en_index++;
149
150 if (result->ob_refcnt == 1) {
151 Py_INCREF(result);
152 Py_DECREF(PyTuple_GET_ITEM(result, 0));
153 Py_DECREF(PyTuple_GET_ITEM(result, 1));
154 } else {
155 result = PyTuple_New(2);
156 if (result == NULL) {
157 Py_DECREF(next_index);
158 Py_DECREF(next_item);
159 return NULL;
160 }
161 }
162 PyTuple_SET_ITEM(result, 0, next_index);
163 PyTuple_SET_ITEM(result, 1, next_item);
164 return result;
165 }
166
167 PyDoc_STRVAR(enum_doc,
168 "enumerate(iterable[, start]) -> iterator for index, value of iterable\n"
169 "\n"
170 "Return an enumerate object. iterable must be another object that supports\n"
171 "iteration. The enumerate object yields pairs containing a count (from\n"
172 "start, which defaults to zero) and a value yielded by the iterable argument.\n"
173 "enumerate is useful for obtaining an indexed list:\n"
174 " (0, seq[0]), (1, seq[1]), (2, seq[2]), ...");
175
176 PyTypeObject PyEnum_Type = {
177 PyVarObject_HEAD_INIT(&PyType_Type, 0)
178 "enumerate", /* tp_name */
179 sizeof(enumobject), /* tp_basicsize */
180 0, /* tp_itemsize */
181 /* methods */
182 (destructor)enum_dealloc, /* tp_dealloc */
183 0, /* tp_print */
184 0, /* tp_getattr */
185 0, /* tp_setattr */
186 0, /* tp_compare */
187 0, /* tp_repr */
188 0, /* tp_as_number */
189 0, /* tp_as_sequence */
190 0, /* tp_as_mapping */
191 0, /* tp_hash */
192 0, /* tp_call */
193 0, /* tp_str */
194 PyObject_GenericGetAttr, /* tp_getattro */
195 0, /* tp_setattro */
196 0, /* tp_as_buffer */
197 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
198 Py_TPFLAGS_BASETYPE, /* tp_flags */
199 enum_doc, /* tp_doc */
200 (traverseproc)enum_traverse, /* tp_traverse */
201 0, /* tp_clear */
202 0, /* tp_richcompare */
203 0, /* tp_weaklistoffset */
204 PyObject_SelfIter, /* tp_iter */
205 (iternextfunc)enum_next, /* tp_iternext */
206 0, /* tp_methods */
207 0, /* tp_members */
208 0, /* tp_getset */
209 0, /* tp_base */
210 0, /* tp_dict */
211 0, /* tp_descr_get */
212 0, /* tp_descr_set */
213 0, /* tp_dictoffset */
214 0, /* tp_init */
215 PyType_GenericAlloc, /* tp_alloc */
216 enum_new, /* tp_new */
217 PyObject_GC_Del, /* tp_free */
218 };
219
220 /* Reversed Object ***************************************************************/
221
222 typedef struct {
223 PyObject_HEAD
224 Py_ssize_t index;
225 PyObject* seq;
226 } reversedobject;
227
228 static PyObject *
reversed_new(PyTypeObject * type,PyObject * args,PyObject * kwds)229 reversed_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
230 {
231 Py_ssize_t n;
232 PyObject *seq, *reversed_meth;
233 static PyObject *reversed_cache = NULL;
234 reversedobject *ro;
235
236 if (type == &PyReversed_Type && !_PyArg_NoKeywords("reversed()", kwds))
237 return NULL;
238
239 if (!PyArg_UnpackTuple(args, "reversed", 1, 1, &seq) )
240 return NULL;
241
242 if (PyInstance_Check(seq)) {
243 reversed_meth = PyObject_GetAttrString(seq, "__reversed__");
244 if (reversed_meth == NULL) {
245 if (PyErr_ExceptionMatches(PyExc_AttributeError))
246 PyErr_Clear();
247 else
248 return NULL;
249 }
250 }
251 else {
252 reversed_meth = _PyObject_LookupSpecial(seq, "__reversed__",
253 &reversed_cache);
254 if (reversed_meth == NULL && PyErr_Occurred())
255 return NULL;
256 }
257 if (reversed_meth != NULL) {
258 PyObject *res = PyObject_CallFunctionObjArgs(reversed_meth, NULL);
259 Py_DECREF(reversed_meth);
260 return res;
261 }
262
263 if (!PySequence_Check(seq)) {
264 PyErr_SetString(PyExc_TypeError,
265 "argument to reversed() must be a sequence");
266 return NULL;
267 }
268
269 n = PySequence_Size(seq);
270 if (n == -1)
271 return NULL;
272
273 ro = (reversedobject *)type->tp_alloc(type, 0);
274 if (ro == NULL)
275 return NULL;
276
277 ro->index = n-1;
278 Py_INCREF(seq);
279 ro->seq = seq;
280 return (PyObject *)ro;
281 }
282
283 static void
reversed_dealloc(reversedobject * ro)284 reversed_dealloc(reversedobject *ro)
285 {
286 PyObject_GC_UnTrack(ro);
287 Py_XDECREF(ro->seq);
288 Py_TYPE(ro)->tp_free(ro);
289 }
290
291 static int
reversed_traverse(reversedobject * ro,visitproc visit,void * arg)292 reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
293 {
294 Py_VISIT(ro->seq);
295 return 0;
296 }
297
298 static PyObject *
reversed_next(reversedobject * ro)299 reversed_next(reversedobject *ro)
300 {
301 PyObject *item;
302 Py_ssize_t index = ro->index;
303
304 if (index >= 0) {
305 item = PySequence_GetItem(ro->seq, index);
306 if (item != NULL) {
307 ro->index--;
308 return item;
309 }
310 if (PyErr_ExceptionMatches(PyExc_IndexError) ||
311 PyErr_ExceptionMatches(PyExc_StopIteration))
312 PyErr_Clear();
313 }
314 ro->index = -1;
315 Py_CLEAR(ro->seq);
316 return NULL;
317 }
318
319 PyDoc_STRVAR(reversed_doc,
320 "reversed(sequence) -> reverse iterator over values of the sequence\n"
321 "\n"
322 "Return a reverse iterator");
323
324 static PyObject *
reversed_len(reversedobject * ro)325 reversed_len(reversedobject *ro)
326 {
327 Py_ssize_t position, seqsize;
328
329 if (ro->seq == NULL)
330 return PyInt_FromLong(0);
331 seqsize = PySequence_Size(ro->seq);
332 if (seqsize == -1)
333 return NULL;
334 position = ro->index + 1;
335 return PyInt_FromSsize_t((seqsize < position) ? 0 : position);
336 }
337
338 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
339
340 static PyMethodDef reversediter_methods[] = {
341 {"__length_hint__", (PyCFunction)reversed_len, METH_NOARGS, length_hint_doc},
342 {NULL, NULL} /* sentinel */
343 };
344
345 PyTypeObject PyReversed_Type = {
346 PyVarObject_HEAD_INIT(&PyType_Type, 0)
347 "reversed", /* tp_name */
348 sizeof(reversedobject), /* tp_basicsize */
349 0, /* tp_itemsize */
350 /* methods */
351 (destructor)reversed_dealloc, /* tp_dealloc */
352 0, /* tp_print */
353 0, /* tp_getattr */
354 0, /* tp_setattr */
355 0, /* tp_compare */
356 0, /* tp_repr */
357 0, /* tp_as_number */
358 0, /* tp_as_sequence */
359 0, /* tp_as_mapping */
360 0, /* tp_hash */
361 0, /* tp_call */
362 0, /* tp_str */
363 PyObject_GenericGetAttr, /* tp_getattro */
364 0, /* tp_setattro */
365 0, /* tp_as_buffer */
366 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
367 Py_TPFLAGS_BASETYPE, /* tp_flags */
368 reversed_doc, /* tp_doc */
369 (traverseproc)reversed_traverse,/* tp_traverse */
370 0, /* tp_clear */
371 0, /* tp_richcompare */
372 0, /* tp_weaklistoffset */
373 PyObject_SelfIter, /* tp_iter */
374 (iternextfunc)reversed_next, /* tp_iternext */
375 reversediter_methods, /* tp_methods */
376 0, /* tp_members */
377 0, /* tp_getset */
378 0, /* tp_base */
379 0, /* tp_dict */
380 0, /* tp_descr_get */
381 0, /* tp_descr_set */
382 0, /* tp_dictoffset */
383 0, /* tp_init */
384 PyType_GenericAlloc, /* tp_alloc */
385 reversed_new, /* tp_new */
386 PyObject_GC_Del, /* tp_free */
387 };
388