1 /* Iterator objects */
2
3 #include "Python.h"
4 #include "pycore_object.h"
5
6 typedef 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
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 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
114 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
115
116 static PyObject *
iter_setstate(seqiterobject * it,PyObject * state)117 iter_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
130 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
131
132 static 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
139 PyTypeObject 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
174 typedef 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
180 PyObject *
PyCallIter_New(PyObject * callable,PyObject * sentinel)181 PyCallIter_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 }
194 static void
calliter_dealloc(calliterobject * it)195 calliter_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
203 static int
calliter_traverse(calliterobject * it,visitproc visit,void * arg)204 calliter_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
211 static PyObject *
calliter_iternext(calliterobject * it)212 calliter_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
243 static PyObject *
calliter_reduce(calliterobject * it,PyObject * Py_UNUSED (ignored))244 calliter_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
253 static PyMethodDef calliter_methods[] = {
254 {"__reduce__", (PyCFunction)calliter_reduce, METH_NOARGS, reduce_doc},
255 {NULL, NULL} /* sentinel */
256 };
257
258 PyTypeObject 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
293 typedef struct {
294 PyObject_HEAD
295 PyObject *wrapped;
296 PyObject *default_value;
297 } anextawaitableobject;
298
299 static void
anextawaitable_dealloc(anextawaitableobject * obj)300 anextawaitable_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
308 static int
anextawaitable_traverse(anextawaitableobject * obj,visitproc visit,void * arg)309 anextawaitable_traverse(anextawaitableobject *obj, visitproc visit, void *arg)
310 {
311 Py_VISIT(obj->wrapped);
312 Py_VISIT(obj->default_value);
313 return 0;
314 }
315
316 static PyObject *
anextawaitable_getiter(anextawaitableobject * obj)317 anextawaitable_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
346 static PyObject *
anextawaitable_iternext(anextawaitableobject * obj)347 anextawaitable_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
385 static PyObject *
anextawaitable_proxy(anextawaitableobject * obj,char * meth,PyObject * arg)386 anextawaitable_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
408 static PyObject *
anextawaitable_send(anextawaitableobject * obj,PyObject * arg)409 anextawaitable_send(anextawaitableobject *obj, PyObject *arg) {
410 return anextawaitable_proxy(obj, "send", arg);
411 }
412
413
414 static PyObject *
anextawaitable_throw(anextawaitableobject * obj,PyObject * arg)415 anextawaitable_throw(anextawaitableobject *obj, PyObject *arg) {
416 return anextawaitable_proxy(obj, "throw", arg);
417 }
418
419
420 static PyObject *
anextawaitable_close(anextawaitableobject * obj,PyObject * arg)421 anextawaitable_close(anextawaitableobject *obj, PyObject *arg) {
422 return anextawaitable_proxy(obj, "close", arg);
423 }
424
425
426 PyDoc_STRVAR(send_doc,
427 "send(arg) -> send 'arg' into the wrapped iterator,\n\
428 return next yielded value or raise StopIteration.");
429
430
431 PyDoc_STRVAR(throw_doc,
432 "throw(typ[,val[,tb]]) -> raise exception in the wrapped iterator,\n\
433 return next yielded value or raise StopIteration.");
434
435
436 PyDoc_STRVAR(close_doc,
437 "close() -> raise GeneratorExit inside generator.");
438
439
440 static 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
448 static PyAsyncMethods anextawaitable_as_async = {
449 PyObject_SelfIter, /* am_await */
450 0, /* am_aiter */
451 0, /* am_anext */
452 0, /* am_send */
453 };
454
455 PyTypeObject _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
487 PyObject *
PyAnextAwaitable_New(PyObject * awaitable,PyObject * default_value)488 PyAnextAwaitable_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