1/* Iterator objects */
2
3#include "Python.h"
4#include "pycore_object.h"
5
6typedef 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
14PyObject *
15PySeqIter_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
33static void
34iter_dealloc(seqiterobject *it)
35{
36 _PyObject_GC_UNTRACK(it);
37 Py_XDECREF(it->it_seq);
38 PyObject_GC_Del(it);
39}
40
41static int
42iter_traverse(seqiterobject *it, visitproc visit, void *arg)
43{
44 Py_VISIT(it->it_seq);
45 return 0;
46}
47
48static PyObject *
49iter_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
81static PyObject *
82iter_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
102PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
103
104static PyObject *
105iter_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
114PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
115
116static PyObject *
117iter_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
130PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
131
132static 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
139PyTypeObject 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
174typedef 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
180PyObject *
181PyCallIter_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}
194static void
195calliter_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
203static int
204calliter_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
211static PyObject *
212calliter_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
243static PyObject *
244calliter_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
253static PyMethodDef calliter_methods[] = {
254 {"__reduce__", (PyCFunction)calliter_reduce, METH_NOARGS, reduce_doc},
255 {NULL, NULL} /* sentinel */
256};
257
258PyTypeObject 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/* -------------------------------------- */
292
293typedef struct {
294 PyObject_HEAD
295 PyObject *wrapped;
296 PyObject *default_value;
297} anextawaitableobject;
298
299static void
300anextawaitable_dealloc(anextawaitableobject *obj)
301{
302 _PyObject_GC_UNTRACK(obj);
303 Py_XDECREF(obj->wrapped);
304 Py_XDECREF(obj->default_value);
305 PyObject_GC_Del(obj);
306}
307
308static int
309anextawaitable_traverse(anextawaitableobject *obj, visitproc visit, void *arg)
310{
311 Py_VISIT(obj->wrapped);
312 Py_VISIT(obj->default_value);
313 return 0;
314}
315
316static PyObject *
317anextawaitable_getiter(anextawaitableobject *obj)
318{
319 assert(obj->wrapped != NULL);
320 PyObject *awaitable = _PyCoro_GetAwaitableIter(obj->wrapped);
321 if (awaitable == NULL) {
322 return NULL;
323 }
324 if (Py_TYPE(awaitable)->tp_iternext == NULL) {
325 /* _PyCoro_GetAwaitableIter returns a Coroutine, a Generator,
326 * or an iterator. Of these, only coroutines lack tp_iternext.
327 */
328 assert(PyCoro_CheckExact(awaitable));
329 unaryfunc getter = Py_TYPE(awaitable)->tp_as_async->am_await;
330 PyObject *new_awaitable = getter(awaitable);
331 if (new_awaitable == NULL) {
332 Py_DECREF(awaitable);
333 return NULL;
334 }
335 Py_SETREF(awaitable, new_awaitable);
336 if (!PyIter_Check(awaitable)) {
337 PyErr_SetString(PyExc_TypeError,
338 "__await__ returned a non-iterable");
339 Py_DECREF(awaitable);
340 return NULL;
341 }
342 }
343 return awaitable;
344}
345
346static PyObject *
347anextawaitable_iternext(anextawaitableobject *obj)
348{
349 /* Consider the following class:
350 *
351 * class A:
352 * async def __anext__(self):
353 * ...
354 * a = A()
355 *
356 * Then `await anext(a)` should call
357 * a.__anext__().__await__().__next__()
358 *
359 * On the other hand, given
360 *
361 * async def agen():
362 * yield 1
363 * yield 2
364 * gen = agen()
365 *
366 * Then `await anext(gen)` can just call
367 * gen.__anext__().__next__()
368 */
369 PyObject *awaitable = anextawaitable_getiter(obj);
370 if (awaitable == NULL) {
371 return NULL;
372 }
373 PyObject *result = (*Py_TYPE(awaitable)->tp_iternext)(awaitable);
374 Py_DECREF(awaitable);
375 if (result != NULL) {
376 return result;
377 }
378 if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
379 _PyGen_SetStopIterationValue(obj->default_value);
380 }
381 return NULL;
382}
383
384
385static PyObject *
386anextawaitable_proxy(anextawaitableobject *obj, char *meth, PyObject *arg) {
387 PyObject *awaitable = anextawaitable_getiter(obj);
388 if (awaitable == NULL) {
389 return NULL;
390 }
391 PyObject *ret = PyObject_CallMethod(awaitable, meth, "O", arg);
392 Py_DECREF(awaitable);
393 if (ret != NULL) {
394 return ret;
395 }
396 if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
397 /* `anextawaitableobject` is only used by `anext()` when
398 * a default value is provided. So when we have a StopAsyncIteration
399 * exception we replace it with a `StopIteration(default)`, as if
400 * it was the return value of `__anext__()` coroutine.
401 */
402 _PyGen_SetStopIterationValue(obj->default_value);
403 }
404 return NULL;
405}
406
407
408static PyObject *
409anextawaitable_send(anextawaitableobject *obj, PyObject *arg) {
410 return anextawaitable_proxy(obj, "send", arg);
411}
412
413
414static PyObject *
415anextawaitable_throw(anextawaitableobject *obj, PyObject *arg) {
416 return anextawaitable_proxy(obj, "throw", arg);
417}
418
419
420static PyObject *
421anextawaitable_close(anextawaitableobject *obj, PyObject *arg) {
422 return anextawaitable_proxy(obj, "close", arg);
423}
424
425
426PyDoc_STRVAR(send_doc,
427"send(arg) -> send 'arg' into the wrapped iterator,\n\
428return next yielded value or raise StopIteration.");
429
430
431PyDoc_STRVAR(throw_doc,
432"throw(typ[,val[,tb]]) -> raise exception in the wrapped iterator,\n\
433return next yielded value or raise StopIteration.");
434
435
436PyDoc_STRVAR(close_doc,
437"close() -> raise GeneratorExit inside generator.");
438
439
440static PyMethodDef anextawaitable_methods[] = {
441 {"send",(PyCFunction)anextawaitable_send, METH_O, send_doc},
442 {"throw",(PyCFunction)anextawaitable_throw, METH_VARARGS, throw_doc},
443 {"close",(PyCFunction)anextawaitable_close, METH_VARARGS, close_doc},
444 {NULL, NULL} /* Sentinel */
445};
446
447
448static PyAsyncMethods anextawaitable_as_async = {
449 PyObject_SelfIter, /* am_await */
450 0, /* am_aiter */
451 0, /* am_anext */
452 0, /* am_send */
453};
454
455PyTypeObject _PyAnextAwaitable_Type = {
456 PyVarObject_HEAD_INIT(&PyType_Type, 0)
457 "anext_awaitable", /* tp_name */
458 sizeof(anextawaitableobject), /* tp_basicsize */
459 0, /* tp_itemsize */
460 /* methods */
461 (destructor)anextawaitable_dealloc, /* tp_dealloc */
462 0, /* tp_vectorcall_offset */
463 0, /* tp_getattr */
464 0, /* tp_setattr */
465 &anextawaitable_as_async, /* tp_as_async */
466 0, /* tp_repr */
467 0, /* tp_as_number */
468 0, /* tp_as_sequence */
469 0, /* tp_as_mapping */
470 0, /* tp_hash */
471 0, /* tp_call */
472 0, /* tp_str */
473 PyObject_GenericGetAttr, /* tp_getattro */
474 0, /* tp_setattro */
475 0, /* tp_as_buffer */
476 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
477 0, /* tp_doc */
478 (traverseproc)anextawaitable_traverse, /* tp_traverse */
479 0, /* tp_clear */
480 0, /* tp_richcompare */
481 0, /* tp_weaklistoffset */
482 PyObject_SelfIter, /* tp_iter */
483 (unaryfunc)anextawaitable_iternext, /* tp_iternext */
484 anextawaitable_methods, /* tp_methods */
485};
486
487PyObject *
488PyAnextAwaitable_New(PyObject *awaitable, PyObject *default_value)
489{
490 anextawaitableobject *anext = PyObject_GC_New(
491 anextawaitableobject, &_PyAnextAwaitable_Type);
492 if (anext == NULL) {
493 return NULL;
494 }
495 Py_INCREF(awaitable);
496 anext->wrapped = awaitable;
497 Py_INCREF(default_value);
498 anext->default_value = default_value;
499 _PyObject_GC_TRACK(anext);
500 return (PyObject *)anext;
501}
502