• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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