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