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