blob: c458cfe73dd06cd843e500a1100019819d98f20f [file] [log] [blame]
Guido van Rossum7dab2422002-04-26 19:40:56 +00001/* enumerate object */
2
3#include "Python.h"
4
5typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00006 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 */
Guido van Rossum7dab2422002-04-26 19:40:56 +000011} enumobject;
12
Guido van Rossum7dab2422002-04-26 19:40:56 +000013static PyObject *
14enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
15{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000016 enumobject *en;
17 PyObject *seq = NULL;
18 PyObject *start = NULL;
19 static char *kwlist[] = {"iterable", "start", 0};
Guido van Rossum7dab2422002-04-26 19:40:56 +000020
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000021 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:enumerate", kwlist,
22 &seq, &start))
23 return NULL;
Guido van Rossum7dab2422002-04-26 19:40:56 +000024
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000025 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(PyLong_Check(start));
35 en->en_index = PyLong_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;
Guido van Rossum7dab2422002-04-26 19:40:56 +000059}
60
61static void
62enum_dealloc(enumobject *en)
63{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000064 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);
Guido van Rossum7dab2422002-04-26 19:40:56 +000069}
70
71static int
72enum_traverse(enumobject *en, visitproc visit, void *arg)
73{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000074 Py_VISIT(en->en_sit);
75 Py_VISIT(en->en_result);
76 Py_VISIT(en->en_longindex);
77 return 0;
Guido van Rossum7dab2422002-04-26 19:40:56 +000078}
79
80static PyObject *
Guido van Rossum8ce8a782007-11-01 19:42:39 +000081enum_next_long(enumobject *en, PyObject* next_item)
82{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000083 static PyObject *one = NULL;
84 PyObject *result = en->en_result;
85 PyObject *next_index;
86 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +000087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000088 if (en->en_longindex == NULL) {
89 en->en_longindex = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
90 if (en->en_longindex == NULL)
91 return NULL;
92 }
93 if (one == NULL) {
94 one = PyLong_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;
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000104
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000105 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;
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000120}
121
122static PyObject *
Guido van Rossum7dab2422002-04-26 19:40:56 +0000123enum_next(enumobject *en)
124{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000125 PyObject *next_index;
126 PyObject *next_item;
127 PyObject *result = en->en_result;
128 PyObject *it = en->en_sit;
Guido van Rossum7dab2422002-04-26 19:40:56 +0000129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000130 next_item = (*Py_TYPE(it)->tp_iternext)(it);
131 if (next_item == NULL)
132 return NULL;
Guido van Rossumca5ed5b2002-07-16 21:02:42 +0000133
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000134 if (en->en_index == PY_SSIZE_T_MAX)
135 return enum_next_long(en, next_item);
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000136
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000137 next_index = PyLong_FromSsize_t(en->en_index);
138 if (next_index == NULL) {
139 Py_DECREF(next_item);
140 return NULL;
141 }
142 en->en_index++;
Raymond Hettingere8b0f042003-05-28 14:05:34 +0000143
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000144 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;
Guido van Rossum7dab2422002-04-26 19:40:56 +0000159}
160
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000161static PyObject *
162enum_reduce(enumobject *en)
163{
164 if (en->en_longindex != NULL)
165 return Py_BuildValue("O(OO)", Py_TYPE(en), en->en_sit, en->en_longindex);
166 else
167 return Py_BuildValue("O(On)", Py_TYPE(en), en->en_sit, en->en_index);
168}
169
170PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
171
172static PyMethodDef enum_methods[] = {
173 {"__reduce__", (PyCFunction)enum_reduce, METH_NOARGS, reduce_doc},
174 {NULL, NULL} /* sentinel */
175};
176
Martin v. Löwis14f8b4c2002-06-13 20:33:02 +0000177PyDoc_STRVAR(enum_doc,
Benjamin Peterson3e5cd1d2010-06-27 21:45:24 +0000178"enumerate(iterable[, start]) -> iterator for index, value of iterable\n"
Jeremy Hyltonfbbe3472003-04-21 20:26:25 +0000179"\n"
Ezio Melottied610932009-09-25 16:12:33 +0000180"Return an enumerate object. iterable must be another object that supports\n"
Jeremy Hyltonfbbe3472003-04-21 20:26:25 +0000181"iteration. The enumerate object yields pairs containing a count (from\n"
Benjamin Peterson3e5cd1d2010-06-27 21:45:24 +0000182"start, which defaults to zero) and a value yielded by the iterable argument.\n"
183"enumerate is useful for obtaining an indexed list:\n"
184" (0, seq[0]), (1, seq[1]), (2, seq[2]), ...");
Guido van Rossum7dab2422002-04-26 19:40:56 +0000185
186PyTypeObject PyEnum_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000187 PyVarObject_HEAD_INIT(&PyType_Type, 0)
188 "enumerate", /* tp_name */
189 sizeof(enumobject), /* tp_basicsize */
190 0, /* tp_itemsize */
191 /* methods */
192 (destructor)enum_dealloc, /* tp_dealloc */
193 0, /* tp_print */
194 0, /* tp_getattr */
195 0, /* tp_setattr */
196 0, /* tp_reserved */
197 0, /* tp_repr */
198 0, /* tp_as_number */
199 0, /* tp_as_sequence */
200 0, /* tp_as_mapping */
201 0, /* tp_hash */
202 0, /* tp_call */
203 0, /* tp_str */
204 PyObject_GenericGetAttr, /* tp_getattro */
205 0, /* tp_setattro */
206 0, /* tp_as_buffer */
207 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
208 Py_TPFLAGS_BASETYPE, /* tp_flags */
209 enum_doc, /* tp_doc */
210 (traverseproc)enum_traverse, /* tp_traverse */
211 0, /* tp_clear */
212 0, /* tp_richcompare */
213 0, /* tp_weaklistoffset */
214 PyObject_SelfIter, /* tp_iter */
215 (iternextfunc)enum_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000216 enum_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000217 0, /* tp_members */
218 0, /* tp_getset */
219 0, /* tp_base */
220 0, /* tp_dict */
221 0, /* tp_descr_get */
222 0, /* tp_descr_set */
223 0, /* tp_dictoffset */
224 0, /* tp_init */
225 PyType_GenericAlloc, /* tp_alloc */
226 enum_new, /* tp_new */
227 PyObject_GC_Del, /* tp_free */
Guido van Rossum7dab2422002-04-26 19:40:56 +0000228};
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000229
230/* Reversed Object ***************************************************************/
231
232typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000233 PyObject_HEAD
234 Py_ssize_t index;
235 PyObject* seq;
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000236} reversedobject;
237
238static PyObject *
239reversed_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
240{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000241 Py_ssize_t n;
242 PyObject *seq, *reversed_meth;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000243 reversedobject *ro;
Benjamin Petersonce798522012-01-22 11:24:29 -0500244 _Py_IDENTIFIER(__reversed__);
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000245
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000246 if (type == &PyReversed_Type && !_PyArg_NoKeywords("reversed()", kwds))
247 return NULL;
Alexandre Vassalottibee32532008-05-16 18:15:12 +0000248
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000249 if (!PyArg_UnpackTuple(args, "reversed", 1, 1, &seq) )
250 return NULL;
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000251
Benjamin Petersonce798522012-01-22 11:24:29 -0500252 reversed_meth = _PyObject_LookupSpecial(seq, &PyId___reversed__);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000253 if (reversed_meth != NULL) {
254 PyObject *res = PyObject_CallFunctionObjArgs(reversed_meth, NULL);
255 Py_DECREF(reversed_meth);
256 return res;
257 }
258 else if (PyErr_Occurred())
259 return NULL;
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000260
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000261 if (!PySequence_Check(seq)) {
262 PyErr_SetString(PyExc_TypeError,
263 "argument to reversed() must be a sequence");
264 return NULL;
265 }
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000267 n = PySequence_Size(seq);
268 if (n == -1)
269 return NULL;
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000270
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000271 ro = (reversedobject *)type->tp_alloc(type, 0);
272 if (ro == NULL)
273 return NULL;
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000274
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000275 ro->index = n-1;
276 Py_INCREF(seq);
277 ro->seq = seq;
278 return (PyObject *)ro;
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000279}
280
281static void
282reversed_dealloc(reversedobject *ro)
283{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000284 PyObject_GC_UnTrack(ro);
285 Py_XDECREF(ro->seq);
286 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000287}
288
289static int
290reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
291{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000292 Py_VISIT(ro->seq);
293 return 0;
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000294}
295
296static PyObject *
297reversed_next(reversedobject *ro)
298{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000299 PyObject *item;
300 Py_ssize_t index = ro->index;
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000301
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000302 if (index >= 0) {
303 item = PySequence_GetItem(ro->seq, index);
304 if (item != NULL) {
305 ro->index--;
306 return item;
307 }
308 if (PyErr_ExceptionMatches(PyExc_IndexError) ||
309 PyErr_ExceptionMatches(PyExc_StopIteration))
310 PyErr_Clear();
311 }
312 ro->index = -1;
313 Py_CLEAR(ro->seq);
314 return NULL;
Raymond Hettinger029dba52004-02-10 09:33:39 +0000315}
316
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000317PyDoc_STRVAR(reversed_doc,
Raymond Hettingerd2afee42004-08-25 19:42:12 +0000318"reversed(sequence) -> reverse iterator over values of the sequence\n"
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000319"\n"
320"Return a reverse iterator");
321
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000322static PyObject *
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000323reversed_len(reversedobject *ro)
324{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000325 Py_ssize_t position, seqsize;
Raymond Hettinger7892b1c2004-04-12 18:10:01 +0000326
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000327 if (ro->seq == NULL)
328 return PyLong_FromLong(0);
329 seqsize = PySequence_Size(ro->seq);
330 if (seqsize == -1)
331 return NULL;
332 position = ro->index + 1;
333 return PyLong_FromSsize_t((seqsize < position) ? 0 : position);
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000334}
335
Armin Rigof5b3e362006-02-11 21:32:43 +0000336PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000337
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000338static PyObject *
339reversed_reduce(reversedobject *ro)
340{
341 if (ro->seq)
342 return Py_BuildValue("O(O)n", Py_TYPE(ro), ro->seq, ro->index);
343 else
344 return Py_BuildValue("O(())", Py_TYPE(ro));
345}
346
347static PyObject *
348reversed_setstate(reversedobject *ro, PyObject *state)
349{
350 Py_ssize_t index = PyLong_AsSsize_t(state);
351 if (index == -1 && PyErr_Occurred())
352 return NULL;
353 if (ro->seq != 0) {
354 Py_ssize_t n = PySequence_Size(ro->seq);
355 if (n < 0)
356 return NULL;
357 if (index < -1)
358 index = -1;
359 else if (index > n-1)
360 index = n-1;
361 ro->index = index;
362 }
363 Py_RETURN_NONE;
364}
365
366PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
367
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000368static PyMethodDef reversediter_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000369 {"__length_hint__", (PyCFunction)reversed_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000370 {"__reduce__", (PyCFunction)reversed_reduce, METH_NOARGS, reduce_doc},
371 {"__setstate__", (PyCFunction)reversed_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000372 {NULL, NULL} /* sentinel */
Raymond Hettinger06353f72004-02-08 10:49:42 +0000373};
374
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000375PyTypeObject PyReversed_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000376 PyVarObject_HEAD_INIT(&PyType_Type, 0)
377 "reversed", /* tp_name */
378 sizeof(reversedobject), /* tp_basicsize */
379 0, /* tp_itemsize */
380 /* methods */
381 (destructor)reversed_dealloc, /* tp_dealloc */
382 0, /* tp_print */
383 0, /* tp_getattr */
384 0, /* tp_setattr */
385 0, /* tp_reserved */
386 0, /* tp_repr */
387 0, /* tp_as_number */
388 0, /* tp_as_sequence */
389 0, /* tp_as_mapping */
390 0, /* tp_hash */
391 0, /* tp_call */
392 0, /* tp_str */
393 PyObject_GenericGetAttr, /* tp_getattro */
394 0, /* tp_setattro */
395 0, /* tp_as_buffer */
396 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
397 Py_TPFLAGS_BASETYPE, /* tp_flags */
398 reversed_doc, /* tp_doc */
399 (traverseproc)reversed_traverse,/* tp_traverse */
400 0, /* tp_clear */
401 0, /* tp_richcompare */
402 0, /* tp_weaklistoffset */
403 PyObject_SelfIter, /* tp_iter */
404 (iternextfunc)reversed_next, /* tp_iternext */
405 reversediter_methods, /* tp_methods */
406 0, /* tp_members */
407 0, /* tp_getset */
408 0, /* tp_base */
409 0, /* tp_dict */
410 0, /* tp_descr_get */
411 0, /* tp_descr_set */
412 0, /* tp_dictoffset */
413 0, /* tp_init */
414 PyType_GenericAlloc, /* tp_alloc */
415 reversed_new, /* tp_new */
416 PyObject_GC_Del, /* tp_free */
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000417};