1 /* Iterator objects */
2
3 #include "Python.h"
4 #include "pycore_object.h"
5 #include "pycore_pymem.h"
6 #include "pycore_pystate.h"
7
8 typedef struct {
9 PyObject_HEAD
10 Py_ssize_t it_index;
11 PyObject *it_seq; /* Set to NULL when iterator is exhausted */
12 } seqiterobject;
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 _Py_IDENTIFIER(iter);
108 if (it->it_seq != NULL)
109 return Py_BuildValue("N(O)n", _PyEval_GetBuiltinId(&PyId_iter),
110 it->it_seq, it->it_index);
111 else
112 return Py_BuildValue("N(())", _PyEval_GetBuiltinId(&PyId_iter));
113 }
114
115 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
116
117 static PyObject *
iter_setstate(seqiterobject * it,PyObject * state)118 iter_setstate(seqiterobject *it, PyObject *state)
119 {
120 Py_ssize_t index = PyLong_AsSsize_t(state);
121 if (index == -1 && PyErr_Occurred())
122 return NULL;
123 if (it->it_seq != NULL) {
124 if (index < 0)
125 index = 0;
126 it->it_index = index;
127 }
128 Py_RETURN_NONE;
129 }
130
131 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
132
133 static PyMethodDef seqiter_methods[] = {
134 {"__length_hint__", (PyCFunction)iter_len, METH_NOARGS, length_hint_doc},
135 {"__reduce__", (PyCFunction)iter_reduce, METH_NOARGS, reduce_doc},
136 {"__setstate__", (PyCFunction)iter_setstate, METH_O, setstate_doc},
137 {NULL, NULL} /* sentinel */
138 };
139
140 PyTypeObject PySeqIter_Type = {
141 PyVarObject_HEAD_INIT(&PyType_Type, 0)
142 "iterator", /* tp_name */
143 sizeof(seqiterobject), /* tp_basicsize */
144 0, /* tp_itemsize */
145 /* methods */
146 (destructor)iter_dealloc, /* tp_dealloc */
147 0, /* tp_vectorcall_offset */
148 0, /* tp_getattr */
149 0, /* tp_setattr */
150 0, /* tp_as_async */
151 0, /* tp_repr */
152 0, /* tp_as_number */
153 0, /* tp_as_sequence */
154 0, /* tp_as_mapping */
155 0, /* tp_hash */
156 0, /* tp_call */
157 0, /* tp_str */
158 PyObject_GenericGetAttr, /* tp_getattro */
159 0, /* tp_setattro */
160 0, /* tp_as_buffer */
161 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
162 0, /* tp_doc */
163 (traverseproc)iter_traverse, /* tp_traverse */
164 0, /* tp_clear */
165 0, /* tp_richcompare */
166 0, /* tp_weaklistoffset */
167 PyObject_SelfIter, /* tp_iter */
168 iter_iternext, /* tp_iternext */
169 seqiter_methods, /* tp_methods */
170 0, /* tp_members */
171 };
172
173 /* -------------------------------------- */
174
175 typedef struct {
176 PyObject_HEAD
177 PyObject *it_callable; /* Set to NULL when iterator is exhausted */
178 PyObject *it_sentinel; /* Set to NULL when iterator is exhausted */
179 } calliterobject;
180
181 PyObject *
PyCallIter_New(PyObject * callable,PyObject * sentinel)182 PyCallIter_New(PyObject *callable, PyObject *sentinel)
183 {
184 calliterobject *it;
185 it = PyObject_GC_New(calliterobject, &PyCallIter_Type);
186 if (it == NULL)
187 return NULL;
188 Py_INCREF(callable);
189 it->it_callable = callable;
190 Py_INCREF(sentinel);
191 it->it_sentinel = sentinel;
192 _PyObject_GC_TRACK(it);
193 return (PyObject *)it;
194 }
195 static void
calliter_dealloc(calliterobject * it)196 calliter_dealloc(calliterobject *it)
197 {
198 _PyObject_GC_UNTRACK(it);
199 Py_XDECREF(it->it_callable);
200 Py_XDECREF(it->it_sentinel);
201 PyObject_GC_Del(it);
202 }
203
204 static int
calliter_traverse(calliterobject * it,visitproc visit,void * arg)205 calliter_traverse(calliterobject *it, visitproc visit, void *arg)
206 {
207 Py_VISIT(it->it_callable);
208 Py_VISIT(it->it_sentinel);
209 return 0;
210 }
211
212 static PyObject *
calliter_iternext(calliterobject * it)213 calliter_iternext(calliterobject *it)
214 {
215 PyObject *result;
216
217 if (it->it_callable == NULL) {
218 return NULL;
219 }
220
221 result = _PyObject_CallNoArg(it->it_callable);
222 if (result != NULL) {
223 int ok;
224
225 ok = PyObject_RichCompareBool(it->it_sentinel, result, Py_EQ);
226 if (ok == 0) {
227 return result; /* Common case, fast path */
228 }
229
230 Py_DECREF(result);
231 if (ok > 0) {
232 Py_CLEAR(it->it_callable);
233 Py_CLEAR(it->it_sentinel);
234 }
235 }
236 else if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
237 PyErr_Clear();
238 Py_CLEAR(it->it_callable);
239 Py_CLEAR(it->it_sentinel);
240 }
241 return NULL;
242 }
243
244 static PyObject *
calliter_reduce(calliterobject * it,PyObject * Py_UNUSED (ignored))245 calliter_reduce(calliterobject *it, PyObject *Py_UNUSED(ignored))
246 {
247 _Py_IDENTIFIER(iter);
248 if (it->it_callable != NULL && it->it_sentinel != NULL)
249 return Py_BuildValue("N(OO)", _PyEval_GetBuiltinId(&PyId_iter),
250 it->it_callable, it->it_sentinel);
251 else
252 return Py_BuildValue("N(())", _PyEval_GetBuiltinId(&PyId_iter));
253 }
254
255 static PyMethodDef calliter_methods[] = {
256 {"__reduce__", (PyCFunction)calliter_reduce, METH_NOARGS, reduce_doc},
257 {NULL, NULL} /* sentinel */
258 };
259
260 PyTypeObject PyCallIter_Type = {
261 PyVarObject_HEAD_INIT(&PyType_Type, 0)
262 "callable_iterator", /* tp_name */
263 sizeof(calliterobject), /* tp_basicsize */
264 0, /* tp_itemsize */
265 /* methods */
266 (destructor)calliter_dealloc, /* tp_dealloc */
267 0, /* tp_vectorcall_offset */
268 0, /* tp_getattr */
269 0, /* tp_setattr */
270 0, /* tp_as_async */
271 0, /* tp_repr */
272 0, /* tp_as_number */
273 0, /* tp_as_sequence */
274 0, /* tp_as_mapping */
275 0, /* tp_hash */
276 0, /* tp_call */
277 0, /* tp_str */
278 PyObject_GenericGetAttr, /* tp_getattro */
279 0, /* tp_setattro */
280 0, /* tp_as_buffer */
281 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
282 0, /* tp_doc */
283 (traverseproc)calliter_traverse, /* tp_traverse */
284 0, /* tp_clear */
285 0, /* tp_richcompare */
286 0, /* tp_weaklistoffset */
287 PyObject_SelfIter, /* tp_iter */
288 (iternextfunc)calliter_iternext, /* tp_iternext */
289 calliter_methods, /* tp_methods */
290 };
291
292
293