1 /* Iterator objects */
2 
3 #include "Python.h"
4 
5 typedef struct {
6     PyObject_HEAD
7     Py_ssize_t it_index;
8     PyObject *it_seq; /* Set to NULL when iterator is exhausted */
9 } seqiterobject;
10 
11 PyObject *
12 PySeqIter_New(PyObject *seq)
13 {
14     seqiterobject *it;
15 
16     if (!PySequence_Check(seq)) {
17         PyErr_BadInternalCall();
18         return NULL;
19     }
20     it = PyObject_GC_New(seqiterobject, &PySeqIter_Type);
21     if (it == NULL)
22         return NULL;
23     it->it_index = 0;
24     Py_INCREF(seq);
25     it->it_seq = seq;
26     _PyObject_GC_TRACK(it);
27     return (PyObject *)it;
28 }
29 
30 static void
31 iter_dealloc(seqiterobject *it)
32 {
33     _PyObject_GC_UNTRACK(it);
34     Py_XDECREF(it->it_seq);
35     PyObject_GC_Del(it);
36 }
37 
38 static int
39 iter_traverse(seqiterobject *it, visitproc visit, void *arg)
40 {
41     Py_VISIT(it->it_seq);
42     return 0;
43 }
44 
45 static PyObject *
46 iter_iternext(PyObject *iterator)
47 {
48     seqiterobject *it;
49     PyObject *seq;
50     PyObject *result;
51 
52     assert(PySeqIter_Check(iterator));
53     it = (seqiterobject *)iterator;
54     seq = it->it_seq;
55     if (seq == NULL)
56         return NULL;
57     if (it->it_index == PY_SSIZE_T_MAX) {
58         PyErr_SetString(PyExc_OverflowError,
59                         "iter index too large");
60         return NULL;
61     }
62 
63     result = PySequence_GetItem(seq, it->it_index);
64     if (result != NULL) {
65         it->it_index++;
66         return result;
67     }
68     if (PyErr_ExceptionMatches(PyExc_IndexError) ||
69         PyErr_ExceptionMatches(PyExc_StopIteration))
70     {
71         PyErr_Clear();
72         it->it_seq = NULL;
73         Py_DECREF(seq);
74     }
75     return NULL;
76 }
77 
78 static PyObject *
79 iter_len(seqiterobject *it)
80 {
81     Py_ssize_t seqsize, len;
82 
83     if (it->it_seq) {
84         if (_PyObject_HasLen(it->it_seq)) {
85             seqsize = PySequence_Size(it->it_seq);
86             if (seqsize == -1)
87                 return NULL;
88         }
89         else {
90             Py_RETURN_NOTIMPLEMENTED;
91         }
92         len = seqsize - it->it_index;
93         if (len >= 0)
94             return PyLong_FromSsize_t(len);
95     }
96     return PyLong_FromLong(0);
97 }
98 
99 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
100 
101 static PyObject *
102 iter_reduce(seqiterobject *it)
103 {
104     if (it->it_seq != NULL)
105         return Py_BuildValue("N(O)n", _PyObject_GetBuiltin("iter"),
106                              it->it_seq, it->it_index);
107     else
108         return Py_BuildValue("N(())", _PyObject_GetBuiltin("iter"));
109 }
110 
111 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
112 
113 static PyObject *
114 iter_setstate(seqiterobject *it, PyObject *state)
115 {
116     Py_ssize_t index = PyLong_AsSsize_t(state);
117     if (index == -1 && PyErr_Occurred())
118         return NULL;
119     if (it->it_seq != NULL) {
120         if (index < 0)
121             index = 0;
122         it->it_index = index;
123     }
124     Py_RETURN_NONE;
125 }
126 
127 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
128 
129 static PyMethodDef seqiter_methods[] = {
130     {"__length_hint__", (PyCFunction)iter_len, METH_NOARGS, length_hint_doc},
131     {"__reduce__", (PyCFunction)iter_reduce, METH_NOARGS, reduce_doc},
132     {"__setstate__", (PyCFunction)iter_setstate, METH_O, setstate_doc},
133     {NULL,              NULL}           /* sentinel */
134 };
135 
136 PyTypeObject PySeqIter_Type = {
137     PyVarObject_HEAD_INIT(&PyType_Type, 0)
138     "iterator",                                 /* tp_name */
139     sizeof(seqiterobject),                      /* tp_basicsize */
140     0,                                          /* tp_itemsize */
141     /* methods */
142     (destructor)iter_dealloc,                   /* tp_dealloc */
143     0,                                          /* tp_print */
144     0,                                          /* tp_getattr */
145     0,                                          /* tp_setattr */
146     0,                                          /* tp_reserved */
147     0,                                          /* tp_repr */
148     0,                                          /* tp_as_number */
149     0,                                          /* tp_as_sequence */
150     0,                                          /* tp_as_mapping */
151     0,                                          /* tp_hash */
152     0,                                          /* tp_call */
153     0,                                          /* tp_str */
154     PyObject_GenericGetAttr,                    /* tp_getattro */
155     0,                                          /* tp_setattro */
156     0,                                          /* tp_as_buffer */
157     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
158     0,                                          /* tp_doc */
159     (traverseproc)iter_traverse,                /* tp_traverse */
160     0,                                          /* tp_clear */
161     0,                                          /* tp_richcompare */
162     0,                                          /* tp_weaklistoffset */
163     PyObject_SelfIter,                          /* tp_iter */
164     iter_iternext,                              /* tp_iternext */
165     seqiter_methods,                            /* tp_methods */
166     0,                                          /* tp_members */
167 };
168 
169 /* -------------------------------------- */
170 
171 typedef struct {
172     PyObject_HEAD
173     PyObject *it_callable; /* Set to NULL when iterator is exhausted */
174     PyObject *it_sentinel; /* Set to NULL when iterator is exhausted */
175 } calliterobject;
176 
177 PyObject *
178 PyCallIter_New(PyObject *callable, PyObject *sentinel)
179 {
180     calliterobject *it;
181     it = PyObject_GC_New(calliterobject, &PyCallIter_Type);
182     if (it == NULL)
183         return NULL;
184     Py_INCREF(callable);
185     it->it_callable = callable;
186     Py_INCREF(sentinel);
187     it->it_sentinel = sentinel;
188     _PyObject_GC_TRACK(it);
189     return (PyObject *)it;
190 }
191 static void
192 calliter_dealloc(calliterobject *it)
193 {
194     _PyObject_GC_UNTRACK(it);
195     Py_XDECREF(it->it_callable);
196     Py_XDECREF(it->it_sentinel);
197     PyObject_GC_Del(it);
198 }
199 
200 static int
201 calliter_traverse(calliterobject *it, visitproc visit, void *arg)
202 {
203     Py_VISIT(it->it_callable);
204     Py_VISIT(it->it_sentinel);
205     return 0;
206 }
207 
208 static PyObject *
209 calliter_iternext(calliterobject *it)
210 {
211     PyObject *result;
212 
213     if (it->it_callable == NULL) {
214         return NULL;
215     }
216 
217     result = _PyObject_CallNoArg(it->it_callable);
218     if (result != NULL) {
219         int ok;
220 
221         ok = PyObject_RichCompareBool(it->it_sentinel, result, Py_EQ);
222         if (ok == 0) {
223             return result; /* Common case, fast path */
224         }
225 
226         Py_DECREF(result);
227         if (ok > 0) {
228             Py_CLEAR(it->it_callable);
229             Py_CLEAR(it->it_sentinel);
230         }
231     }
232     else if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
233         PyErr_Clear();
234         Py_CLEAR(it->it_callable);
235         Py_CLEAR(it->it_sentinel);
236     }
237     return NULL;
238 }
239 
240 static PyObject *
241 calliter_reduce(calliterobject *it)
242 {
243     if (it->it_callable != NULL && it->it_sentinel != NULL)
244         return Py_BuildValue("N(OO)", _PyObject_GetBuiltin("iter"),
245                              it->it_callable, it->it_sentinel);
246     else
247         return Py_BuildValue("N(())", _PyObject_GetBuiltin("iter"));
248 }
249 
250 static PyMethodDef calliter_methods[] = {
251     {"__reduce__", (PyCFunction)calliter_reduce, METH_NOARGS, reduce_doc},
252     {NULL,              NULL}           /* sentinel */
253 };
254 
255 PyTypeObject PyCallIter_Type = {
256     PyVarObject_HEAD_INIT(&PyType_Type, 0)
257     "callable_iterator",                        /* tp_name */
258     sizeof(calliterobject),                     /* tp_basicsize */
259     0,                                          /* tp_itemsize */
260     /* methods */
261     (destructor)calliter_dealloc,               /* tp_dealloc */
262     0,                                          /* tp_print */
263     0,                                          /* tp_getattr */
264     0,                                          /* tp_setattr */
265     0,                                          /* tp_reserved */
266     0,                                          /* tp_repr */
267     0,                                          /* tp_as_number */
268     0,                                          /* tp_as_sequence */
269     0,                                          /* tp_as_mapping */
270     0,                                          /* tp_hash */
271     0,                                          /* tp_call */
272     0,                                          /* tp_str */
273     PyObject_GenericGetAttr,                    /* tp_getattro */
274     0,                                          /* tp_setattro */
275     0,                                          /* tp_as_buffer */
276     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
277     0,                                          /* tp_doc */
278     (traverseproc)calliter_traverse,            /* tp_traverse */
279     0,                                          /* tp_clear */
280     0,                                          /* tp_richcompare */
281     0,                                          /* tp_weaklistoffset */
282     PyObject_SelfIter,                          /* tp_iter */
283     (iternextfunc)calliter_iternext,            /* tp_iternext */
284     calliter_methods,                           /* tp_methods */
285 };
286 
287 
288