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