1 /* Iterator objects */
2
3 #include "Python.h"
4 #include "pycore_call.h" // _PyObject_CallNoArgs()
5 #include "pycore_object.h" // _PyObject_GC_TRACK()
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 *
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 it->it_seq = Py_NewRef(seq);
27 _PyObject_GC_TRACK(it);
28 return (PyObject *)it;
29 }
30
31 static void
32 iter_dealloc(seqiterobject *it)
33 {
34 _PyObject_GC_UNTRACK(it);
35 Py_XDECREF(it->it_seq);
36 PyObject_GC_Del(it);
37 }
38
39 static int
40 iter_traverse(seqiterobject *it, visitproc visit, void *arg)
41 {
42 Py_VISIT(it->it_seq);
43 return 0;
44 }
45
46 static PyObject *
47 iter_iternext(PyObject *iterator)
48 {
49 seqiterobject *it;
50 PyObject *seq;
51 PyObject *result;
52
53 assert(PySeqIter_Check(iterator));
54 it = (seqiterobject *)iterator;
55 seq = it->it_seq;
56 if (seq == NULL)
57 return NULL;
58 if (it->it_index == PY_SSIZE_T_MAX) {
59 PyErr_SetString(PyExc_OverflowError,
60 "iter index too large");
61 return NULL;
62 }
63
64 result = PySequence_GetItem(seq, it->it_index);
65 if (result != NULL) {
66 it->it_index++;
67 return result;
68 }
69 if (PyErr_ExceptionMatches(PyExc_IndexError) ||
70 PyErr_ExceptionMatches(PyExc_StopIteration))
71 {
72 PyErr_Clear();
73 it->it_seq = NULL;
74 Py_DECREF(seq);
75 }
76 return NULL;
77 }
78
79 static PyObject *
80 iter_len(seqiterobject *it, PyObject *Py_UNUSED(ignored))
81 {
82 Py_ssize_t seqsize, len;
83
84 if (it->it_seq) {
85 if (_PyObject_HasLen(it->it_seq)) {
86 seqsize = PySequence_Size(it->it_seq);
87 if (seqsize == -1)
88 return NULL;
89 }
90 else {
91 Py_RETURN_NOTIMPLEMENTED;
92 }
93 len = seqsize - it->it_index;
94 if (len >= 0)
95 return PyLong_FromSsize_t(len);
96 }
97 return PyLong_FromLong(0);
98 }
99
100 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
101
102 static PyObject *
103 iter_reduce(seqiterobject *it, PyObject *Py_UNUSED(ignored))
104 {
105 PyObject *iter = _PyEval_GetBuiltin(&_Py_ID(iter));
106
107 /* _PyEval_GetBuiltin can invoke arbitrary code,
108 * call must be before access of iterator pointers.
109 * see issue #101765 */
110
111 if (it->it_seq != NULL)
112 return Py_BuildValue("N(O)n", iter, it->it_seq, it->it_index);
113 else
114 return Py_BuildValue("N(())", iter);
115 }
116
117 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
118
119 static PyObject *
120 iter_setstate(seqiterobject *it, PyObject *state)
121 {
122 Py_ssize_t index = PyLong_AsSsize_t(state);
123 if (index == -1 && PyErr_Occurred())
124 return NULL;
125 if (it->it_seq != NULL) {
126 if (index < 0)
127 index = 0;
128 it->it_index = index;
129 }
130 Py_RETURN_NONE;
131 }
132
133 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
134
135 static PyMethodDef seqiter_methods[] = {
136 {"__length_hint__", (PyCFunction)iter_len, METH_NOARGS, length_hint_doc},
137 {"__reduce__", (PyCFunction)iter_reduce, METH_NOARGS, reduce_doc},
138 {"__setstate__", (PyCFunction)iter_setstate, METH_O, setstate_doc},
139 {NULL, NULL} /* sentinel */
140 };
141
142 PyTypeObject PySeqIter_Type = {
143 PyVarObject_HEAD_INIT(&PyType_Type, 0)
144 "iterator", /* tp_name */
145 sizeof(seqiterobject), /* tp_basicsize */
146 0, /* tp_itemsize */
147 /* methods */
148 (destructor)iter_dealloc, /* tp_dealloc */
149 0, /* tp_vectorcall_offset */
150 0, /* tp_getattr */
151 0, /* tp_setattr */
152 0, /* tp_as_async */
153 0, /* tp_repr */
154 0, /* tp_as_number */
155 0, /* tp_as_sequence */
156 0, /* tp_as_mapping */
157 0, /* tp_hash */
158 0, /* tp_call */
159 0, /* tp_str */
160 PyObject_GenericGetAttr, /* tp_getattro */
161 0, /* tp_setattro */
162 0, /* tp_as_buffer */
163 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
164 0, /* tp_doc */
165 (traverseproc)iter_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 iter_iternext, /* tp_iternext */
171 seqiter_methods, /* tp_methods */
172 0, /* tp_members */
173 };
174
175 /* -------------------------------------- */
176
177 typedef struct {
178 PyObject_HEAD
179 PyObject *it_callable; /* Set to NULL when iterator is exhausted */
180 PyObject *it_sentinel; /* Set to NULL when iterator is exhausted */
181 } calliterobject;
182
183 PyObject *
184 PyCallIter_New(PyObject *callable, PyObject *sentinel)
185 {
186 calliterobject *it;
187 it = PyObject_GC_New(calliterobject, &PyCallIter_Type);
188 if (it == NULL)
189 return NULL;
190 it->it_callable = Py_NewRef(callable);
191 it->it_sentinel = Py_NewRef(sentinel);
192 _PyObject_GC_TRACK(it);
193 return (PyObject *)it;
194 }
195 static void
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
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 *
213 calliter_iternext(calliterobject *it)
214 {
215 PyObject *result;
216
217 if (it->it_callable == NULL) {
218 return NULL;
219 }
220
221 result = _PyObject_CallNoArgs(it->it_callable);
222 if (result != NULL && it->it_sentinel != 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 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 Py_XDECREF(result);
241 return NULL;
242 }
243
244 static PyObject *
245 calliter_reduce(calliterobject *it, PyObject *Py_UNUSED(ignored))
246 {
247 PyObject *iter = _PyEval_GetBuiltin(&_Py_ID(iter));
248
249 /* _PyEval_GetBuiltin can invoke arbitrary code,
250 * call must be before access of iterator pointers.
251 * see issue #101765 */
252
253 if (it->it_callable != NULL && it->it_sentinel != NULL)
254 return Py_BuildValue("N(OO)", iter, it->it_callable, it->it_sentinel);
255 else
256 return Py_BuildValue("N(())", iter);
257 }
258
259 static PyMethodDef calliter_methods[] = {
260 {"__reduce__", (PyCFunction)calliter_reduce, METH_NOARGS, reduce_doc},
261 {NULL, NULL} /* sentinel */
262 };
263
264 PyTypeObject PyCallIter_Type = {
265 PyVarObject_HEAD_INIT(&PyType_Type, 0)
266 "callable_iterator", /* tp_name */
267 sizeof(calliterobject), /* tp_basicsize */
268 0, /* tp_itemsize */
269 /* methods */
270 (destructor)calliter_dealloc, /* tp_dealloc */
271 0, /* tp_vectorcall_offset */
272 0, /* tp_getattr */
273 0, /* tp_setattr */
274 0, /* tp_as_async */
275 0, /* tp_repr */
276 0, /* tp_as_number */
277 0, /* tp_as_sequence */
278 0, /* tp_as_mapping */
279 0, /* tp_hash */
280 0, /* tp_call */
281 0, /* tp_str */
282 PyObject_GenericGetAttr, /* tp_getattro */
283 0, /* tp_setattro */
284 0, /* tp_as_buffer */
285 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
286 0, /* tp_doc */
287 (traverseproc)calliter_traverse, /* tp_traverse */
288 0, /* tp_clear */
289 0, /* tp_richcompare */
290 0, /* tp_weaklistoffset */
291 PyObject_SelfIter, /* tp_iter */
292 (iternextfunc)calliter_iternext, /* tp_iternext */
293 calliter_methods, /* tp_methods */
294 };
295
296
297 /* -------------------------------------- */
298
299 typedef struct {
300 PyObject_HEAD
301 PyObject *wrapped;
302 PyObject *default_value;
303 } anextawaitableobject;
304
305 static void
306 anextawaitable_dealloc(anextawaitableobject *obj)
307 {
308 _PyObject_GC_UNTRACK(obj);
309 Py_XDECREF(obj->wrapped);
310 Py_XDECREF(obj->default_value);
311 PyObject_GC_Del(obj);
312 }
313
314 static int
315 anextawaitable_traverse(anextawaitableobject *obj, visitproc visit, void *arg)
316 {
317 Py_VISIT(obj->wrapped);
318 Py_VISIT(obj->default_value);
319 return 0;
320 }
321
322 static PyObject *
323 anextawaitable_getiter(anextawaitableobject *obj)
324 {
325 assert(obj->wrapped != NULL);
326 PyObject *awaitable = _PyCoro_GetAwaitableIter(obj->wrapped);
327 if (awaitable == NULL) {
328 return NULL;
329 }
330 if (Py_TYPE(awaitable)->tp_iternext == NULL) {
331 /* _PyCoro_GetAwaitableIter returns a Coroutine, a Generator,
332 * or an iterator. Of these, only coroutines lack tp_iternext.
333 */
334 assert(PyCoro_CheckExact(awaitable));
335 unaryfunc getter = Py_TYPE(awaitable)->tp_as_async->am_await;
336 PyObject *new_awaitable = getter(awaitable);
337 if (new_awaitable == NULL) {
338 Py_DECREF(awaitable);
339 return NULL;
340 }
341 Py_SETREF(awaitable, new_awaitable);
342 if (!PyIter_Check(awaitable)) {
343 PyErr_SetString(PyExc_TypeError,
344 "__await__ returned a non-iterable");
345 Py_DECREF(awaitable);
346 return NULL;
347 }
348 }
349 return awaitable;
350 }
351
352 static PyObject *
353 anextawaitable_iternext(anextawaitableobject *obj)
354 {
355 /* Consider the following class:
356 *
357 * class A:
358 * async def __anext__(self):
359 * ...
360 * a = A()
361 *
362 * Then `await anext(a)` should call
363 * a.__anext__().__await__().__next__()
364 *
365 * On the other hand, given
366 *
367 * async def agen():
368 * yield 1
369 * yield 2
370 * gen = agen()
371 *
372 * Then `await anext(gen)` can just call
373 * gen.__anext__().__next__()
374 */
375 PyObject *awaitable = anextawaitable_getiter(obj);
376 if (awaitable == NULL) {
377 return NULL;
378 }
379 PyObject *result = (*Py_TYPE(awaitable)->tp_iternext)(awaitable);
380 Py_DECREF(awaitable);
381 if (result != NULL) {
382 return result;
383 }
384 if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
385 _PyGen_SetStopIterationValue(obj->default_value);
386 }
387 return NULL;
388 }
389
390
391 static PyObject *
392 anextawaitable_proxy(anextawaitableobject *obj, char *meth, PyObject *arg) {
393 PyObject *awaitable = anextawaitable_getiter(obj);
394 if (awaitable == NULL) {
395 return NULL;
396 }
397 PyObject *ret = PyObject_CallMethod(awaitable, meth, "O", arg);
398 Py_DECREF(awaitable);
399 if (ret != NULL) {
400 return ret;
401 }
402 if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
403 /* `anextawaitableobject` is only used by `anext()` when
404 * a default value is provided. So when we have a StopAsyncIteration
405 * exception we replace it with a `StopIteration(default)`, as if
406 * it was the return value of `__anext__()` coroutine.
407 */
408 _PyGen_SetStopIterationValue(obj->default_value);
409 }
410 return NULL;
411 }
412
413
414 static PyObject *
415 anextawaitable_send(anextawaitableobject *obj, PyObject *arg) {
416 return anextawaitable_proxy(obj, "send", arg);
417 }
418
419
420 static PyObject *
421 anextawaitable_throw(anextawaitableobject *obj, PyObject *arg) {
422 return anextawaitable_proxy(obj, "throw", arg);
423 }
424
425
426 static PyObject *
427 anextawaitable_close(anextawaitableobject *obj, PyObject *arg) {
428 return anextawaitable_proxy(obj, "close", arg);
429 }
430
431
432 PyDoc_STRVAR(send_doc,
433 "send(arg) -> send 'arg' into the wrapped iterator,\n\
434 return next yielded value or raise StopIteration.");
435
436
437 PyDoc_STRVAR(throw_doc,
438 "throw(value)\n\
439 throw(typ[,val[,tb]])\n\
440 \n\
441 raise exception in the wrapped iterator, return next yielded value\n\
442 or raise StopIteration.\n\
443 the (type, val, tb) signature is deprecated, \n\
444 and may be removed in a future version of Python.");
445
446
447 PyDoc_STRVAR(close_doc,
448 "close() -> raise GeneratorExit inside generator.");
449
450
451 static PyMethodDef anextawaitable_methods[] = {
452 {"send",(PyCFunction)anextawaitable_send, METH_O, send_doc},
453 {"throw",(PyCFunction)anextawaitable_throw, METH_VARARGS, throw_doc},
454 {"close",(PyCFunction)anextawaitable_close, METH_VARARGS, close_doc},
455 {NULL, NULL} /* Sentinel */
456 };
457
458
459 static PyAsyncMethods anextawaitable_as_async = {
460 PyObject_SelfIter, /* am_await */
461 0, /* am_aiter */
462 0, /* am_anext */
463 0, /* am_send */
464 };
465
466 PyTypeObject _PyAnextAwaitable_Type = {
467 PyVarObject_HEAD_INIT(&PyType_Type, 0)
468 "anext_awaitable", /* tp_name */
469 sizeof(anextawaitableobject), /* tp_basicsize */
470 0, /* tp_itemsize */
471 /* methods */
472 (destructor)anextawaitable_dealloc, /* tp_dealloc */
473 0, /* tp_vectorcall_offset */
474 0, /* tp_getattr */
475 0, /* tp_setattr */
476 &anextawaitable_as_async, /* tp_as_async */
477 0, /* tp_repr */
478 0, /* tp_as_number */
479 0, /* tp_as_sequence */
480 0, /* tp_as_mapping */
481 0, /* tp_hash */
482 0, /* tp_call */
483 0, /* tp_str */
484 PyObject_GenericGetAttr, /* tp_getattro */
485 0, /* tp_setattro */
486 0, /* tp_as_buffer */
487 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
488 0, /* tp_doc */
489 (traverseproc)anextawaitable_traverse, /* tp_traverse */
490 0, /* tp_clear */
491 0, /* tp_richcompare */
492 0, /* tp_weaklistoffset */
493 PyObject_SelfIter, /* tp_iter */
494 (unaryfunc)anextawaitable_iternext, /* tp_iternext */
495 anextawaitable_methods, /* tp_methods */
496 };
497
498 PyObject *
499 PyAnextAwaitable_New(PyObject *awaitable, PyObject *default_value)
500 {
501 anextawaitableobject *anext = PyObject_GC_New(
502 anextawaitableobject, &_PyAnextAwaitable_Type);
503 if (anext == NULL) {
504 return NULL;
505 }
506 anext->wrapped = Py_NewRef(awaitable);
507 anext->default_value = Py_NewRef(default_value);
508 _PyObject_GC_TRACK(anext);
509 return (PyObject *)anext;
510 }