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