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