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