1 2 #define PY_SSIZE_T_CLEAN 3 #include "Python.h" 4 #include "structmember.h" 5 6 /* Itertools module written and maintained 7 by Raymond D. Hettinger <python (at) rcn.com> 8 */ 9 10 11 /* groupby object ************************************************************/ 12 13 typedef struct { 14 PyObject_HEAD 15 PyObject *it; 16 PyObject *keyfunc; 17 PyObject *tgtkey; 18 PyObject *currkey; 19 PyObject *currvalue; 20 } groupbyobject; 21 22 static PyTypeObject groupby_type; 23 static PyObject *_grouper_create(groupbyobject *, PyObject *); 24 25 static PyObject * 26 groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 27 { 28 static char *kwargs[] = {"iterable", "key", NULL}; 29 groupbyobject *gbo; 30 PyObject *it, *keyfunc = Py_None; 31 32 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs, 33 &it, &keyfunc)) 34 return NULL; 35 36 gbo = (groupbyobject *)type->tp_alloc(type, 0); 37 if (gbo == NULL) 38 return NULL; 39 gbo->tgtkey = NULL; 40 gbo->currkey = NULL; 41 gbo->currvalue = NULL; 42 gbo->keyfunc = keyfunc; 43 Py_INCREF(keyfunc); 44 gbo->it = PyObject_GetIter(it); 45 if (gbo->it == NULL) { 46 Py_DECREF(gbo); 47 return NULL; 48 } 49 return (PyObject *)gbo; 50 } 51 52 static void 53 groupby_dealloc(groupbyobject *gbo) 54 { 55 PyObject_GC_UnTrack(gbo); 56 Py_XDECREF(gbo->it); 57 Py_XDECREF(gbo->keyfunc); 58 Py_XDECREF(gbo->tgtkey); 59 Py_XDECREF(gbo->currkey); 60 Py_XDECREF(gbo->currvalue); 61 Py_TYPE(gbo)->tp_free(gbo); 62 } 63 64 static int 65 groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg) 66 { 67 Py_VISIT(gbo->it); 68 Py_VISIT(gbo->keyfunc); 69 Py_VISIT(gbo->tgtkey); 70 Py_VISIT(gbo->currkey); 71 Py_VISIT(gbo->currvalue); 72 return 0; 73 } 74 75 static PyObject * 76 groupby_next(groupbyobject *gbo) 77 { 78 PyObject *newvalue, *newkey, *r, *grouper; 79 80 /* skip to next iteration group */ 81 for (;;) { 82 if (gbo->currkey == NULL) 83 /* pass */; 84 else if (gbo->tgtkey == NULL) 85 break; 86 else { 87 int rcmp; 88 89 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ); 90 if (rcmp == -1) 91 return NULL; 92 else if (rcmp == 0) 93 break; 94 } 95 96 newvalue = PyIter_Next(gbo->it); 97 if (newvalue == NULL) 98 return NULL; 99 100 if (gbo->keyfunc == Py_None) { 101 newkey = newvalue; 102 Py_INCREF(newvalue); 103 } else { 104 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL); 105 if (newkey == NULL) { 106 Py_DECREF(newvalue); 107 return NULL; 108 } 109 } 110 111 Py_XSETREF(gbo->currkey, newkey); 112 Py_XSETREF(gbo->currvalue, newvalue); 113 } 114 115 Py_INCREF(gbo->currkey); 116 Py_XSETREF(gbo->tgtkey, gbo->currkey); 117 118 grouper = _grouper_create(gbo, gbo->tgtkey); 119 if (grouper == NULL) 120 return NULL; 121 122 r = PyTuple_Pack(2, gbo->currkey, grouper); 123 Py_DECREF(grouper); 124 return r; 125 } 126 127 static PyObject * 128 groupby_reduce(groupbyobject *lz) 129 { 130 /* reduce as a 'new' call with an optional 'setstate' if groupby 131 * has started 132 */ 133 PyObject *value; 134 if (lz->tgtkey && lz->currkey && lz->currvalue) 135 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz), 136 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey); 137 else 138 value = Py_BuildValue("O(OO)", Py_TYPE(lz), 139 lz->it, lz->keyfunc); 140 141 return value; 142 } 143 144 PyDoc_STRVAR(reduce_doc, "Return state information for pickling."); 145 146 static PyObject * 147 groupby_setstate(groupbyobject *lz, PyObject *state) 148 { 149 PyObject *currkey, *currvalue, *tgtkey; 150 if (!PyTuple_Check(state)) { 151 PyErr_SetString(PyExc_TypeError, "state is not a tuple"); 152 return NULL; 153 } 154 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) { 155 return NULL; 156 } 157 Py_INCREF(currkey); 158 Py_XSETREF(lz->currkey, currkey); 159 Py_INCREF(currvalue); 160 Py_XSETREF(lz->currvalue, currvalue); 161 Py_INCREF(tgtkey); 162 Py_XSETREF(lz->tgtkey, tgtkey); 163 Py_RETURN_NONE; 164 } 165 166 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling."); 167 168 static PyMethodDef groupby_methods[] = { 169 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS, 170 reduce_doc}, 171 {"__setstate__", (PyCFunction)groupby_setstate, METH_O, 172 setstate_doc}, 173 {NULL, NULL} /* sentinel */ 174 }; 175 176 PyDoc_STRVAR(groupby_doc, 177 "groupby(iterable[, keyfunc]) -> create an iterator which returns\n\ 178 (key, sub-iterator) grouped by each value of key(value).\n"); 179 180 static PyTypeObject groupby_type = { 181 PyVarObject_HEAD_INIT(NULL, 0) 182 "itertools.groupby", /* tp_name */ 183 sizeof(groupbyobject), /* tp_basicsize */ 184 0, /* tp_itemsize */ 185 /* methods */ 186 (destructor)groupby_dealloc, /* tp_dealloc */ 187 0, /* tp_print */ 188 0, /* tp_getattr */ 189 0, /* tp_setattr */ 190 0, /* tp_reserved */ 191 0, /* tp_repr */ 192 0, /* tp_as_number */ 193 0, /* tp_as_sequence */ 194 0, /* tp_as_mapping */ 195 0, /* tp_hash */ 196 0, /* tp_call */ 197 0, /* tp_str */ 198 PyObject_GenericGetAttr, /* tp_getattro */ 199 0, /* tp_setattro */ 200 0, /* tp_as_buffer */ 201 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 202 Py_TPFLAGS_BASETYPE, /* tp_flags */ 203 groupby_doc, /* tp_doc */ 204 (traverseproc)groupby_traverse, /* tp_traverse */ 205 0, /* tp_clear */ 206 0, /* tp_richcompare */ 207 0, /* tp_weaklistoffset */ 208 PyObject_SelfIter, /* tp_iter */ 209 (iternextfunc)groupby_next, /* tp_iternext */ 210 groupby_methods, /* tp_methods */ 211 0, /* tp_members */ 212 0, /* tp_getset */ 213 0, /* tp_base */ 214 0, /* tp_dict */ 215 0, /* tp_descr_get */ 216 0, /* tp_descr_set */ 217 0, /* tp_dictoffset */ 218 0, /* tp_init */ 219 0, /* tp_alloc */ 220 groupby_new, /* tp_new */ 221 PyObject_GC_Del, /* tp_free */ 222 }; 223 224 225 /* _grouper object (internal) ************************************************/ 226 227 typedef struct { 228 PyObject_HEAD 229 PyObject *parent; 230 PyObject *tgtkey; 231 } _grouperobject; 232 233 static PyTypeObject _grouper_type; 234 235 static PyObject * 236 _grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 237 { 238 PyObject *parent, *tgtkey; 239 240 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey)) 241 return NULL; 242 243 return _grouper_create((groupbyobject*) parent, tgtkey); 244 } 245 246 static PyObject * 247 _grouper_create(groupbyobject *parent, PyObject *tgtkey) 248 { 249 _grouperobject *igo; 250 251 igo = PyObject_GC_New(_grouperobject, &_grouper_type); 252 if (igo == NULL) 253 return NULL; 254 igo->parent = (PyObject *)parent; 255 Py_INCREF(parent); 256 igo->tgtkey = tgtkey; 257 Py_INCREF(tgtkey); 258 259 PyObject_GC_Track(igo); 260 return (PyObject *)igo; 261 } 262 263 static void 264 _grouper_dealloc(_grouperobject *igo) 265 { 266 PyObject_GC_UnTrack(igo); 267 Py_DECREF(igo->parent); 268 Py_DECREF(igo->tgtkey); 269 PyObject_GC_Del(igo); 270 } 271 272 static int 273 _grouper_traverse(_grouperobject *igo, visitproc visit, void *arg) 274 { 275 Py_VISIT(igo->parent); 276 Py_VISIT(igo->tgtkey); 277 return 0; 278 } 279 280 static PyObject * 281 _grouper_next(_grouperobject *igo) 282 { 283 groupbyobject *gbo = (groupbyobject *)igo->parent; 284 PyObject *newvalue, *newkey, *r; 285 int rcmp; 286 287 if (gbo->currvalue == NULL) { 288 newvalue = PyIter_Next(gbo->it); 289 if (newvalue == NULL) 290 return NULL; 291 292 if (gbo->keyfunc == Py_None) { 293 newkey = newvalue; 294 Py_INCREF(newvalue); 295 } else { 296 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL); 297 if (newkey == NULL) { 298 Py_DECREF(newvalue); 299 return NULL; 300 } 301 } 302 303 assert(gbo->currkey == NULL); 304 gbo->currkey = newkey; 305 gbo->currvalue = newvalue; 306 } 307 308 assert(gbo->currkey != NULL); 309 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ); 310 if (rcmp <= 0) 311 /* got any error or current group is end */ 312 return NULL; 313 314 r = gbo->currvalue; 315 gbo->currvalue = NULL; 316 Py_CLEAR(gbo->currkey); 317 318 return r; 319 } 320 321 static PyObject * 322 _grouper_reduce(_grouperobject *lz) 323 { 324 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey); 325 } 326 327 static PyMethodDef _grouper_methods[] = { 328 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS, 329 reduce_doc}, 330 {NULL, NULL} /* sentinel */ 331 }; 332 333 334 static PyTypeObject _grouper_type = { 335 PyVarObject_HEAD_INIT(NULL, 0) 336 "itertools._grouper", /* tp_name */ 337 sizeof(_grouperobject), /* tp_basicsize */ 338 0, /* tp_itemsize */ 339 /* methods */ 340 (destructor)_grouper_dealloc, /* tp_dealloc */ 341 0, /* tp_print */ 342 0, /* tp_getattr */ 343 0, /* tp_setattr */ 344 0, /* tp_reserved */ 345 0, /* tp_repr */ 346 0, /* tp_as_number */ 347 0, /* tp_as_sequence */ 348 0, /* tp_as_mapping */ 349 0, /* tp_hash */ 350 0, /* tp_call */ 351 0, /* tp_str */ 352 PyObject_GenericGetAttr, /* tp_getattro */ 353 0, /* tp_setattro */ 354 0, /* tp_as_buffer */ 355 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */ 356 0, /* tp_doc */ 357 (traverseproc)_grouper_traverse, /* tp_traverse */ 358 0, /* tp_clear */ 359 0, /* tp_richcompare */ 360 0, /* tp_weaklistoffset */ 361 PyObject_SelfIter, /* tp_iter */ 362 (iternextfunc)_grouper_next, /* tp_iternext */ 363 _grouper_methods, /* tp_methods */ 364 0, /* tp_members */ 365 0, /* tp_getset */ 366 0, /* tp_base */ 367 0, /* tp_dict */ 368 0, /* tp_descr_get */ 369 0, /* tp_descr_set */ 370 0, /* tp_dictoffset */ 371 0, /* tp_init */ 372 0, /* tp_alloc */ 373 _grouper_new, /* tp_new */ 374 PyObject_GC_Del, /* tp_free */ 375 }; 376 377 378 /* tee object and with supporting function and objects ***********************/ 379 380 /* The teedataobject pre-allocates space for LINKCELLS number of objects. 381 To help the object fit neatly inside cache lines (space for 16 to 32 382 pointers), the value should be a multiple of 16 minus space for 383 the other structure members including PyHEAD overhead. The larger the 384 value, the less memory overhead per object and the less time spent 385 allocating/deallocating new links. The smaller the number, the less 386 wasted space and the more rapid freeing of older data. 387 */ 388 #define LINKCELLS 57 389 390 typedef struct { 391 PyObject_HEAD 392 PyObject *it; 393 int numread; /* 0 <= numread <= LINKCELLS */ 394 PyObject *nextlink; 395 PyObject *(values[LINKCELLS]); 396 } teedataobject; 397 398 typedef struct { 399 PyObject_HEAD 400 teedataobject *dataobj; 401 int index; /* 0 <= index <= LINKCELLS */ 402 PyObject *weakreflist; 403 } teeobject; 404 405 static PyTypeObject teedataobject_type; 406 407 static PyObject * 408 teedataobject_newinternal(PyObject *it) 409 { 410 teedataobject *tdo; 411 412 tdo = PyObject_GC_New(teedataobject, &teedataobject_type); 413 if (tdo == NULL) 414 return NULL; 415 416 tdo->numread = 0; 417 tdo->nextlink = NULL; 418 Py_INCREF(it); 419 tdo->it = it; 420 PyObject_GC_Track(tdo); 421 return (PyObject *)tdo; 422 } 423 424 static PyObject * 425 teedataobject_jumplink(teedataobject *tdo) 426 { 427 if (tdo->nextlink == NULL) 428 tdo->nextlink = teedataobject_newinternal(tdo->it); 429 Py_XINCREF(tdo->nextlink); 430 return tdo->nextlink; 431 } 432 433 static PyObject * 434 teedataobject_getitem(teedataobject *tdo, int i) 435 { 436 PyObject *value; 437 438 assert(i < LINKCELLS); 439 if (i < tdo->numread) 440 value = tdo->values[i]; 441 else { 442 /* this is the lead iterator, so fetch more data */ 443 assert(i == tdo->numread); 444 value = PyIter_Next(tdo->it); 445 if (value == NULL) 446 return NULL; 447 tdo->numread++; 448 tdo->values[i] = value; 449 } 450 Py_INCREF(value); 451 return value; 452 } 453 454 static int 455 teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg) 456 { 457 int i; 458 459 Py_VISIT(tdo->it); 460 for (i = 0; i < tdo->numread; i++) 461 Py_VISIT(tdo->values[i]); 462 Py_VISIT(tdo->nextlink); 463 return 0; 464 } 465 466 static void 467 teedataobject_safe_decref(PyObject *obj) 468 { 469 while (obj && Py_TYPE(obj) == &teedataobject_type && 470 Py_REFCNT(obj) == 1) { 471 PyObject *nextlink = ((teedataobject *)obj)->nextlink; 472 ((teedataobject *)obj)->nextlink = NULL; 473 Py_DECREF(obj); 474 obj = nextlink; 475 } 476 Py_XDECREF(obj); 477 } 478 479 static int 480 teedataobject_clear(teedataobject *tdo) 481 { 482 int i; 483 PyObject *tmp; 484 485 Py_CLEAR(tdo->it); 486 for (i=0 ; i<tdo->numread ; i++) 487 Py_CLEAR(tdo->values[i]); 488 tmp = tdo->nextlink; 489 tdo->nextlink = NULL; 490 teedataobject_safe_decref(tmp); 491 return 0; 492 } 493 494 static void 495 teedataobject_dealloc(teedataobject *tdo) 496 { 497 PyObject_GC_UnTrack(tdo); 498 teedataobject_clear(tdo); 499 PyObject_GC_Del(tdo); 500 } 501 502 static PyObject * 503 teedataobject_reduce(teedataobject *tdo) 504 { 505 int i; 506 /* create a temporary list of already iterated values */ 507 PyObject *values = PyList_New(tdo->numread); 508 509 if (!values) 510 return NULL; 511 for (i=0 ; i<tdo->numread ; i++) { 512 Py_INCREF(tdo->values[i]); 513 PyList_SET_ITEM(values, i, tdo->values[i]); 514 } 515 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it, 516 values, 517 tdo->nextlink ? tdo->nextlink : Py_None); 518 } 519 520 static PyTypeObject teedataobject_type; 521 522 static PyObject * 523 teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw) 524 { 525 teedataobject *tdo; 526 PyObject *it, *values, *next; 527 Py_ssize_t i, len; 528 529 assert(type == &teedataobject_type); 530 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next)) 531 return NULL; 532 533 tdo = (teedataobject *)teedataobject_newinternal(it); 534 if (!tdo) 535 return NULL; 536 537 len = PyList_GET_SIZE(values); 538 if (len > LINKCELLS) 539 goto err; 540 for (i=0; i<len; i++) { 541 tdo->values[i] = PyList_GET_ITEM(values, i); 542 Py_INCREF(tdo->values[i]); 543 } 544 /* len <= LINKCELLS < INT_MAX */ 545 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int); 546 547 if (len == LINKCELLS) { 548 if (next != Py_None) { 549 if (Py_TYPE(next) != &teedataobject_type) 550 goto err; 551 assert(tdo->nextlink == NULL); 552 Py_INCREF(next); 553 tdo->nextlink = next; 554 } 555 } else { 556 if (next != Py_None) 557 goto err; /* shouldn't have a next if we are not full */ 558 } 559 return (PyObject*)tdo; 560 561 err: 562 Py_XDECREF(tdo); 563 PyErr_SetString(PyExc_ValueError, "Invalid arguments"); 564 return NULL; 565 } 566 567 static PyMethodDef teedataobject_methods[] = { 568 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS, 569 reduce_doc}, 570 {NULL, NULL} /* sentinel */ 571 }; 572 573 PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects."); 574 575 static PyTypeObject teedataobject_type = { 576 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */ 577 "itertools._tee_dataobject", /* tp_name */ 578 sizeof(teedataobject), /* tp_basicsize */ 579 0, /* tp_itemsize */ 580 /* methods */ 581 (destructor)teedataobject_dealloc, /* tp_dealloc */ 582 0, /* tp_print */ 583 0, /* tp_getattr */ 584 0, /* tp_setattr */ 585 0, /* tp_reserved */ 586 0, /* tp_repr */ 587 0, /* tp_as_number */ 588 0, /* tp_as_sequence */ 589 0, /* tp_as_mapping */ 590 0, /* tp_hash */ 591 0, /* tp_call */ 592 0, /* tp_str */ 593 PyObject_GenericGetAttr, /* tp_getattro */ 594 0, /* tp_setattro */ 595 0, /* tp_as_buffer */ 596 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */ 597 teedataobject_doc, /* tp_doc */ 598 (traverseproc)teedataobject_traverse, /* tp_traverse */ 599 (inquiry)teedataobject_clear, /* tp_clear */ 600 0, /* tp_richcompare */ 601 0, /* tp_weaklistoffset */ 602 0, /* tp_iter */ 603 0, /* tp_iternext */ 604 teedataobject_methods, /* tp_methods */ 605 0, /* tp_members */ 606 0, /* tp_getset */ 607 0, /* tp_base */ 608 0, /* tp_dict */ 609 0, /* tp_descr_get */ 610 0, /* tp_descr_set */ 611 0, /* tp_dictoffset */ 612 0, /* tp_init */ 613 0, /* tp_alloc */ 614 teedataobject_new, /* tp_new */ 615 PyObject_GC_Del, /* tp_free */ 616 }; 617 618 619 static PyTypeObject tee_type; 620 621 static PyObject * 622 tee_next(teeobject *to) 623 { 624 PyObject *value, *link; 625 626 if (to->index >= LINKCELLS) { 627 link = teedataobject_jumplink(to->dataobj); 628 if (link == NULL) 629 return NULL; 630 Py_SETREF(to->dataobj, (teedataobject *)link); 631 to->index = 0; 632 } 633 value = teedataobject_getitem(to->dataobj, to->index); 634 if (value == NULL) 635 return NULL; 636 to->index++; 637 return value; 638 } 639 640 static int 641 tee_traverse(teeobject *to, visitproc visit, void *arg) 642 { 643 Py_VISIT((PyObject *)to->dataobj); 644 return 0; 645 } 646 647 static PyObject * 648 tee_copy(teeobject *to) 649 { 650 teeobject *newto; 651 652 newto = PyObject_GC_New(teeobject, &tee_type); 653 if (newto == NULL) 654 return NULL; 655 Py_INCREF(to->dataobj); 656 newto->dataobj = to->dataobj; 657 newto->index = to->index; 658 newto->weakreflist = NULL; 659 PyObject_GC_Track(newto); 660 return (PyObject *)newto; 661 } 662 663 PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator."); 664 665 static PyObject * 666 tee_fromiterable(PyObject *iterable) 667 { 668 teeobject *to; 669 PyObject *it = NULL; 670 671 it = PyObject_GetIter(iterable); 672 if (it == NULL) 673 return NULL; 674 if (PyObject_TypeCheck(it, &tee_type)) { 675 to = (teeobject *)tee_copy((teeobject *)it); 676 goto done; 677 } 678 679 to = PyObject_GC_New(teeobject, &tee_type); 680 if (to == NULL) 681 goto done; 682 to->dataobj = (teedataobject *)teedataobject_newinternal(it); 683 if (!to->dataobj) { 684 PyObject_GC_Del(to); 685 to = NULL; 686 goto done; 687 } 688 689 to->index = 0; 690 to->weakreflist = NULL; 691 PyObject_GC_Track(to); 692 done: 693 Py_XDECREF(it); 694 return (PyObject *)to; 695 } 696 697 static PyObject * 698 tee_new(PyTypeObject *type, PyObject *args, PyObject *kw) 699 { 700 PyObject *iterable; 701 702 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable)) 703 return NULL; 704 return tee_fromiterable(iterable); 705 } 706 707 static int 708 tee_clear(teeobject *to) 709 { 710 if (to->weakreflist != NULL) 711 PyObject_ClearWeakRefs((PyObject *) to); 712 Py_CLEAR(to->dataobj); 713 return 0; 714 } 715 716 static void 717 tee_dealloc(teeobject *to) 718 { 719 PyObject_GC_UnTrack(to); 720 tee_clear(to); 721 PyObject_GC_Del(to); 722 } 723 724 static PyObject * 725 tee_reduce(teeobject *to) 726 { 727 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index); 728 } 729 730 static PyObject * 731 tee_setstate(teeobject *to, PyObject *state) 732 { 733 teedataobject *tdo; 734 int index; 735 if (!PyTuple_Check(state)) { 736 PyErr_SetString(PyExc_TypeError, "state is not a tuple"); 737 return NULL; 738 } 739 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) { 740 return NULL; 741 } 742 if (index < 0 || index > LINKCELLS) { 743 PyErr_SetString(PyExc_ValueError, "Index out of range"); 744 return NULL; 745 } 746 Py_INCREF(tdo); 747 Py_XSETREF(to->dataobj, tdo); 748 to->index = index; 749 Py_RETURN_NONE; 750 } 751 752 PyDoc_STRVAR(teeobject_doc, 753 "Iterator wrapped to make it copyable"); 754 755 static PyMethodDef tee_methods[] = { 756 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc}, 757 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc}, 758 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc}, 759 {NULL, NULL} /* sentinel */ 760 }; 761 762 static PyTypeObject tee_type = { 763 PyVarObject_HEAD_INIT(NULL, 0) 764 "itertools._tee", /* tp_name */ 765 sizeof(teeobject), /* tp_basicsize */ 766 0, /* tp_itemsize */ 767 /* methods */ 768 (destructor)tee_dealloc, /* tp_dealloc */ 769 0, /* tp_print */ 770 0, /* tp_getattr */ 771 0, /* tp_setattr */ 772 0, /* tp_reserved */ 773 0, /* tp_repr */ 774 0, /* tp_as_number */ 775 0, /* tp_as_sequence */ 776 0, /* tp_as_mapping */ 777 0, /* tp_hash */ 778 0, /* tp_call */ 779 0, /* tp_str */ 780 0, /* tp_getattro */ 781 0, /* tp_setattro */ 782 0, /* tp_as_buffer */ 783 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */ 784 teeobject_doc, /* tp_doc */ 785 (traverseproc)tee_traverse, /* tp_traverse */ 786 (inquiry)tee_clear, /* tp_clear */ 787 0, /* tp_richcompare */ 788 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */ 789 PyObject_SelfIter, /* tp_iter */ 790 (iternextfunc)tee_next, /* tp_iternext */ 791 tee_methods, /* tp_methods */ 792 0, /* tp_members */ 793 0, /* tp_getset */ 794 0, /* tp_base */ 795 0, /* tp_dict */ 796 0, /* tp_descr_get */ 797 0, /* tp_descr_set */ 798 0, /* tp_dictoffset */ 799 0, /* tp_init */ 800 0, /* tp_alloc */ 801 tee_new, /* tp_new */ 802 PyObject_GC_Del, /* tp_free */ 803 }; 804 805 static PyObject * 806 tee(PyObject *self, PyObject *args) 807 { 808 Py_ssize_t i, n=2; 809 PyObject *it, *iterable, *copyable, *result; 810 _Py_IDENTIFIER(__copy__); 811 812 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n)) 813 return NULL; 814 if (n < 0) { 815 PyErr_SetString(PyExc_ValueError, "n must be >= 0"); 816 return NULL; 817 } 818 result = PyTuple_New(n); 819 if (result == NULL) 820 return NULL; 821 if (n == 0) 822 return result; 823 it = PyObject_GetIter(iterable); 824 if (it == NULL) { 825 Py_DECREF(result); 826 return NULL; 827 } 828 if (!_PyObject_HasAttrId(it, &PyId___copy__)) { 829 copyable = tee_fromiterable(it); 830 Py_DECREF(it); 831 if (copyable == NULL) { 832 Py_DECREF(result); 833 return NULL; 834 } 835 } else 836 copyable = it; 837 PyTuple_SET_ITEM(result, 0, copyable); 838 for (i=1 ; i<n ; i++) { 839 840 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL); 841 if (copyable == NULL) { 842 Py_DECREF(result); 843 return NULL; 844 } 845 PyTuple_SET_ITEM(result, i, copyable); 846 } 847 return result; 848 } 849 850 PyDoc_STRVAR(tee_doc, 851 "tee(iterable, n=2) --> tuple of n independent iterators."); 852 853 854 /* cycle object **************************************************************/ 855 856 typedef struct { 857 PyObject_HEAD 858 PyObject *it; 859 PyObject *saved; 860 Py_ssize_t index; 861 int firstpass; 862 } cycleobject; 863 864 static PyTypeObject cycle_type; 865 866 static PyObject * 867 cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 868 { 869 PyObject *it; 870 PyObject *iterable; 871 PyObject *saved; 872 cycleobject *lz; 873 874 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds)) 875 return NULL; 876 877 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable)) 878 return NULL; 879 880 /* Get iterator. */ 881 it = PyObject_GetIter(iterable); 882 if (it == NULL) 883 return NULL; 884 885 saved = PyList_New(0); 886 if (saved == NULL) { 887 Py_DECREF(it); 888 return NULL; 889 } 890 891 /* create cycleobject structure */ 892 lz = (cycleobject *)type->tp_alloc(type, 0); 893 if (lz == NULL) { 894 Py_DECREF(it); 895 Py_DECREF(saved); 896 return NULL; 897 } 898 lz->it = it; 899 lz->saved = saved; 900 lz->index = 0; 901 lz->firstpass = 0; 902 903 return (PyObject *)lz; 904 } 905 906 static void 907 cycle_dealloc(cycleobject *lz) 908 { 909 PyObject_GC_UnTrack(lz); 910 Py_XDECREF(lz->it); 911 Py_XDECREF(lz->saved); 912 Py_TYPE(lz)->tp_free(lz); 913 } 914 915 static int 916 cycle_traverse(cycleobject *lz, visitproc visit, void *arg) 917 { 918 if (lz->it) 919 Py_VISIT(lz->it); 920 Py_VISIT(lz->saved); 921 return 0; 922 } 923 924 static PyObject * 925 cycle_next(cycleobject *lz) 926 { 927 PyObject *item; 928 929 if (lz->it != NULL) { 930 item = PyIter_Next(lz->it); 931 if (item != NULL) { 932 if (lz->firstpass) 933 return item; 934 if (PyList_Append(lz->saved, item)) { 935 Py_DECREF(item); 936 return NULL; 937 } 938 return item; 939 } 940 /* Note: StopIteration is already cleared by PyIter_Next() */ 941 if (PyErr_Occurred()) 942 return NULL; 943 Py_CLEAR(lz->it); 944 } 945 if (Py_SIZE(lz->saved) == 0) 946 return NULL; 947 item = PyList_GET_ITEM(lz->saved, lz->index); 948 lz->index++; 949 if (lz->index >= Py_SIZE(lz->saved)) 950 lz->index = 0; 951 Py_INCREF(item); 952 return item; 953 } 954 955 static PyObject * 956 cycle_reduce(cycleobject *lz) 957 { 958 /* Create a new cycle with the iterator tuple, then set the saved state */ 959 if (lz->it == NULL) { 960 PyObject *it = PyObject_GetIter(lz->saved); 961 if (it == NULL) 962 return NULL; 963 if (lz->index != 0) { 964 _Py_IDENTIFIER(__setstate__); 965 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__, 966 "n", lz->index); 967 if (res == NULL) { 968 Py_DECREF(it); 969 return NULL; 970 } 971 Py_DECREF(res); 972 } 973 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1); 974 } 975 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved, 976 lz->firstpass); 977 } 978 979 static PyObject * 980 cycle_setstate(cycleobject *lz, PyObject *state) 981 { 982 PyObject *saved=NULL; 983 int firstpass; 984 if (!PyTuple_Check(state)) { 985 PyErr_SetString(PyExc_TypeError, "state is not a tuple"); 986 return NULL; 987 } 988 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) { 989 return NULL; 990 } 991 Py_INCREF(saved); 992 Py_XSETREF(lz->saved, saved); 993 lz->firstpass = firstpass != 0; 994 lz->index = 0; 995 Py_RETURN_NONE; 996 } 997 998 static PyMethodDef cycle_methods[] = { 999 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS, 1000 reduce_doc}, 1001 {"__setstate__", (PyCFunction)cycle_setstate, METH_O, 1002 setstate_doc}, 1003 {NULL, NULL} /* sentinel */ 1004 }; 1005 1006 PyDoc_STRVAR(cycle_doc, 1007 "cycle(iterable) --> cycle object\n\ 1008 \n\ 1009 Return elements from the iterable until it is exhausted.\n\ 1010 Then repeat the sequence indefinitely."); 1011 1012 static PyTypeObject cycle_type = { 1013 PyVarObject_HEAD_INIT(NULL, 0) 1014 "itertools.cycle", /* tp_name */ 1015 sizeof(cycleobject), /* tp_basicsize */ 1016 0, /* tp_itemsize */ 1017 /* methods */ 1018 (destructor)cycle_dealloc, /* tp_dealloc */ 1019 0, /* tp_print */ 1020 0, /* tp_getattr */ 1021 0, /* tp_setattr */ 1022 0, /* tp_reserved */ 1023 0, /* tp_repr */ 1024 0, /* tp_as_number */ 1025 0, /* tp_as_sequence */ 1026 0, /* tp_as_mapping */ 1027 0, /* tp_hash */ 1028 0, /* tp_call */ 1029 0, /* tp_str */ 1030 PyObject_GenericGetAttr, /* tp_getattro */ 1031 0, /* tp_setattro */ 1032 0, /* tp_as_buffer */ 1033 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 1034 Py_TPFLAGS_BASETYPE, /* tp_flags */ 1035 cycle_doc, /* tp_doc */ 1036 (traverseproc)cycle_traverse, /* tp_traverse */ 1037 0, /* tp_clear */ 1038 0, /* tp_richcompare */ 1039 0, /* tp_weaklistoffset */ 1040 PyObject_SelfIter, /* tp_iter */ 1041 (iternextfunc)cycle_next, /* tp_iternext */ 1042 cycle_methods, /* tp_methods */ 1043 0, /* tp_members */ 1044 0, /* tp_getset */ 1045 0, /* tp_base */ 1046 0, /* tp_dict */ 1047 0, /* tp_descr_get */ 1048 0, /* tp_descr_set */ 1049 0, /* tp_dictoffset */ 1050 0, /* tp_init */ 1051 0, /* tp_alloc */ 1052 cycle_new, /* tp_new */ 1053 PyObject_GC_Del, /* tp_free */ 1054 }; 1055 1056 1057 /* dropwhile object **********************************************************/ 1058 1059 typedef struct { 1060 PyObject_HEAD 1061 PyObject *func; 1062 PyObject *it; 1063 long start; 1064 } dropwhileobject; 1065 1066 static PyTypeObject dropwhile_type; 1067 1068 static PyObject * 1069 dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 1070 { 1071 PyObject *func, *seq; 1072 PyObject *it; 1073 dropwhileobject *lz; 1074 1075 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds)) 1076 return NULL; 1077 1078 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq)) 1079 return NULL; 1080 1081 /* Get iterator. */ 1082 it = PyObject_GetIter(seq); 1083 if (it == NULL) 1084 return NULL; 1085 1086 /* create dropwhileobject structure */ 1087 lz = (dropwhileobject *)type->tp_alloc(type, 0); 1088 if (lz == NULL) { 1089 Py_DECREF(it); 1090 return NULL; 1091 } 1092 Py_INCREF(func); 1093 lz->func = func; 1094 lz->it = it; 1095 lz->start = 0; 1096 1097 return (PyObject *)lz; 1098 } 1099 1100 static void 1101 dropwhile_dealloc(dropwhileobject *lz) 1102 { 1103 PyObject_GC_UnTrack(lz); 1104 Py_XDECREF(lz->func); 1105 Py_XDECREF(lz->it); 1106 Py_TYPE(lz)->tp_free(lz); 1107 } 1108 1109 static int 1110 dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg) 1111 { 1112 Py_VISIT(lz->it); 1113 Py_VISIT(lz->func); 1114 return 0; 1115 } 1116 1117 static PyObject * 1118 dropwhile_next(dropwhileobject *lz) 1119 { 1120 PyObject *item, *good; 1121 PyObject *it = lz->it; 1122 long ok; 1123 PyObject *(*iternext)(PyObject *); 1124 1125 iternext = *Py_TYPE(it)->tp_iternext; 1126 for (;;) { 1127 item = iternext(it); 1128 if (item == NULL) 1129 return NULL; 1130 if (lz->start == 1) 1131 return item; 1132 1133 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL); 1134 if (good == NULL) { 1135 Py_DECREF(item); 1136 return NULL; 1137 } 1138 ok = PyObject_IsTrue(good); 1139 Py_DECREF(good); 1140 if (ok == 0) { 1141 lz->start = 1; 1142 return item; 1143 } 1144 Py_DECREF(item); 1145 if (ok < 0) 1146 return NULL; 1147 } 1148 } 1149 1150 static PyObject * 1151 dropwhile_reduce(dropwhileobject *lz) 1152 { 1153 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start); 1154 } 1155 1156 static PyObject * 1157 dropwhile_setstate(dropwhileobject *lz, PyObject *state) 1158 { 1159 int start = PyObject_IsTrue(state); 1160 if (start < 0) 1161 return NULL; 1162 lz->start = start; 1163 Py_RETURN_NONE; 1164 } 1165 1166 static PyMethodDef dropwhile_methods[] = { 1167 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS, 1168 reduce_doc}, 1169 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O, 1170 setstate_doc}, 1171 {NULL, NULL} /* sentinel */ 1172 }; 1173 1174 PyDoc_STRVAR(dropwhile_doc, 1175 "dropwhile(predicate, iterable) --> dropwhile object\n\ 1176 \n\ 1177 Drop items from the iterable while predicate(item) is true.\n\ 1178 Afterwards, return every element until the iterable is exhausted."); 1179 1180 static PyTypeObject dropwhile_type = { 1181 PyVarObject_HEAD_INIT(NULL, 0) 1182 "itertools.dropwhile", /* tp_name */ 1183 sizeof(dropwhileobject), /* tp_basicsize */ 1184 0, /* tp_itemsize */ 1185 /* methods */ 1186 (destructor)dropwhile_dealloc, /* tp_dealloc */ 1187 0, /* tp_print */ 1188 0, /* tp_getattr */ 1189 0, /* tp_setattr */ 1190 0, /* tp_reserved */ 1191 0, /* tp_repr */ 1192 0, /* tp_as_number */ 1193 0, /* tp_as_sequence */ 1194 0, /* tp_as_mapping */ 1195 0, /* tp_hash */ 1196 0, /* tp_call */ 1197 0, /* tp_str */ 1198 PyObject_GenericGetAttr, /* tp_getattro */ 1199 0, /* tp_setattro */ 1200 0, /* tp_as_buffer */ 1201 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 1202 Py_TPFLAGS_BASETYPE, /* tp_flags */ 1203 dropwhile_doc, /* tp_doc */ 1204 (traverseproc)dropwhile_traverse, /* tp_traverse */ 1205 0, /* tp_clear */ 1206 0, /* tp_richcompare */ 1207 0, /* tp_weaklistoffset */ 1208 PyObject_SelfIter, /* tp_iter */ 1209 (iternextfunc)dropwhile_next, /* tp_iternext */ 1210 dropwhile_methods, /* tp_methods */ 1211 0, /* tp_members */ 1212 0, /* tp_getset */ 1213 0, /* tp_base */ 1214 0, /* tp_dict */ 1215 0, /* tp_descr_get */ 1216 0, /* tp_descr_set */ 1217 0, /* tp_dictoffset */ 1218 0, /* tp_init */ 1219 0, /* tp_alloc */ 1220 dropwhile_new, /* tp_new */ 1221 PyObject_GC_Del, /* tp_free */ 1222 }; 1223 1224 1225 /* takewhile object **********************************************************/ 1226 1227 typedef struct { 1228 PyObject_HEAD 1229 PyObject *func; 1230 PyObject *it; 1231 long stop; 1232 } takewhileobject; 1233 1234 static PyTypeObject takewhile_type; 1235 1236 static PyObject * 1237 takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 1238 { 1239 PyObject *func, *seq; 1240 PyObject *it; 1241 takewhileobject *lz; 1242 1243 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds)) 1244 return NULL; 1245 1246 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq)) 1247 return NULL; 1248 1249 /* Get iterator. */ 1250 it = PyObject_GetIter(seq); 1251 if (it == NULL) 1252 return NULL; 1253 1254 /* create takewhileobject structure */ 1255 lz = (takewhileobject *)type->tp_alloc(type, 0); 1256 if (lz == NULL) { 1257 Py_DECREF(it); 1258 return NULL; 1259 } 1260 Py_INCREF(func); 1261 lz->func = func; 1262 lz->it = it; 1263 lz->stop = 0; 1264 1265 return (PyObject *)lz; 1266 } 1267 1268 static void 1269 takewhile_dealloc(takewhileobject *lz) 1270 { 1271 PyObject_GC_UnTrack(lz); 1272 Py_XDECREF(lz->func); 1273 Py_XDECREF(lz->it); 1274 Py_TYPE(lz)->tp_free(lz); 1275 } 1276 1277 static int 1278 takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg) 1279 { 1280 Py_VISIT(lz->it); 1281 Py_VISIT(lz->func); 1282 return 0; 1283 } 1284 1285 static PyObject * 1286 takewhile_next(takewhileobject *lz) 1287 { 1288 PyObject *item, *good; 1289 PyObject *it = lz->it; 1290 long ok; 1291 1292 if (lz->stop == 1) 1293 return NULL; 1294 1295 item = (*Py_TYPE(it)->tp_iternext)(it); 1296 if (item == NULL) 1297 return NULL; 1298 1299 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL); 1300 if (good == NULL) { 1301 Py_DECREF(item); 1302 return NULL; 1303 } 1304 ok = PyObject_IsTrue(good); 1305 Py_DECREF(good); 1306 if (ok > 0) 1307 return item; 1308 Py_DECREF(item); 1309 if (ok == 0) 1310 lz->stop = 1; 1311 return NULL; 1312 } 1313 1314 static PyObject * 1315 takewhile_reduce(takewhileobject *lz) 1316 { 1317 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop); 1318 } 1319 1320 static PyObject * 1321 takewhile_reduce_setstate(takewhileobject *lz, PyObject *state) 1322 { 1323 int stop = PyObject_IsTrue(state); 1324 1325 if (stop < 0) 1326 return NULL; 1327 lz->stop = stop; 1328 Py_RETURN_NONE; 1329 } 1330 1331 static PyMethodDef takewhile_reduce_methods[] = { 1332 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS, 1333 reduce_doc}, 1334 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O, 1335 setstate_doc}, 1336 {NULL, NULL} /* sentinel */ 1337 }; 1338 PyDoc_STRVAR(takewhile_doc, 1339 "takewhile(predicate, iterable) --> takewhile object\n\ 1340 \n\ 1341 Return successive entries from an iterable as long as the \n\ 1342 predicate evaluates to true for each entry."); 1343 1344 static PyTypeObject takewhile_type = { 1345 PyVarObject_HEAD_INIT(NULL, 0) 1346 "itertools.takewhile", /* tp_name */ 1347 sizeof(takewhileobject), /* tp_basicsize */ 1348 0, /* tp_itemsize */ 1349 /* methods */ 1350 (destructor)takewhile_dealloc, /* tp_dealloc */ 1351 0, /* tp_print */ 1352 0, /* tp_getattr */ 1353 0, /* tp_setattr */ 1354 0, /* tp_reserved */ 1355 0, /* tp_repr */ 1356 0, /* tp_as_number */ 1357 0, /* tp_as_sequence */ 1358 0, /* tp_as_mapping */ 1359 0, /* tp_hash */ 1360 0, /* tp_call */ 1361 0, /* tp_str */ 1362 PyObject_GenericGetAttr, /* tp_getattro */ 1363 0, /* tp_setattro */ 1364 0, /* tp_as_buffer */ 1365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 1366 Py_TPFLAGS_BASETYPE, /* tp_flags */ 1367 takewhile_doc, /* tp_doc */ 1368 (traverseproc)takewhile_traverse, /* tp_traverse */ 1369 0, /* tp_clear */ 1370 0, /* tp_richcompare */ 1371 0, /* tp_weaklistoffset */ 1372 PyObject_SelfIter, /* tp_iter */ 1373 (iternextfunc)takewhile_next, /* tp_iternext */ 1374 takewhile_reduce_methods, /* tp_methods */ 1375 0, /* tp_members */ 1376 0, /* tp_getset */ 1377 0, /* tp_base */ 1378 0, /* tp_dict */ 1379 0, /* tp_descr_get */ 1380 0, /* tp_descr_set */ 1381 0, /* tp_dictoffset */ 1382 0, /* tp_init */ 1383 0, /* tp_alloc */ 1384 takewhile_new, /* tp_new */ 1385 PyObject_GC_Del, /* tp_free */ 1386 }; 1387 1388 1389 /* islice object *************************************************************/ 1390 1391 typedef struct { 1392 PyObject_HEAD 1393 PyObject *it; 1394 Py_ssize_t next; 1395 Py_ssize_t stop; 1396 Py_ssize_t step; 1397 Py_ssize_t cnt; 1398 } isliceobject; 1399 1400 static PyTypeObject islice_type; 1401 1402 static PyObject * 1403 islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 1404 { 1405 PyObject *seq; 1406 Py_ssize_t start=0, stop=-1, step=1; 1407 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL; 1408 Py_ssize_t numargs; 1409 isliceobject *lz; 1410 1411 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds)) 1412 return NULL; 1413 1414 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3)) 1415 return NULL; 1416 1417 numargs = PyTuple_Size(args); 1418 if (numargs == 2) { 1419 if (a1 != Py_None) { 1420 stop = PyLong_AsSsize_t(a1); 1421 if (stop == -1) { 1422 if (PyErr_Occurred()) 1423 PyErr_Clear(); 1424 PyErr_SetString(PyExc_ValueError, 1425 "Stop argument for islice() must be None or " 1426 "an integer: 0 <= x <= sys.maxsize."); 1427 return NULL; 1428 } 1429 } 1430 } else { 1431 if (a1 != Py_None) 1432 start = PyLong_AsSsize_t(a1); 1433 if (start == -1 && PyErr_Occurred()) 1434 PyErr_Clear(); 1435 if (a2 != Py_None) { 1436 stop = PyLong_AsSsize_t(a2); 1437 if (stop == -1) { 1438 if (PyErr_Occurred()) 1439 PyErr_Clear(); 1440 PyErr_SetString(PyExc_ValueError, 1441 "Stop argument for islice() must be None or " 1442 "an integer: 0 <= x <= sys.maxsize."); 1443 return NULL; 1444 } 1445 } 1446 } 1447 if (start<0 || stop<-1) { 1448 PyErr_SetString(PyExc_ValueError, 1449 "Indices for islice() must be None or " 1450 "an integer: 0 <= x <= sys.maxsize."); 1451 return NULL; 1452 } 1453 1454 if (a3 != NULL) { 1455 if (a3 != Py_None) 1456 step = PyLong_AsSsize_t(a3); 1457 if (step == -1 && PyErr_Occurred()) 1458 PyErr_Clear(); 1459 } 1460 if (step<1) { 1461 PyErr_SetString(PyExc_ValueError, 1462 "Step for islice() must be a positive integer or None."); 1463 return NULL; 1464 } 1465 1466 /* Get iterator. */ 1467 it = PyObject_GetIter(seq); 1468 if (it == NULL) 1469 return NULL; 1470 1471 /* create isliceobject structure */ 1472 lz = (isliceobject *)type->tp_alloc(type, 0); 1473 if (lz == NULL) { 1474 Py_DECREF(it); 1475 return NULL; 1476 } 1477 lz->it = it; 1478 lz->next = start; 1479 lz->stop = stop; 1480 lz->step = step; 1481 lz->cnt = 0L; 1482 1483 return (PyObject *)lz; 1484 } 1485 1486 static void 1487 islice_dealloc(isliceobject *lz) 1488 { 1489 PyObject_GC_UnTrack(lz); 1490 Py_XDECREF(lz->it); 1491 Py_TYPE(lz)->tp_free(lz); 1492 } 1493 1494 static int 1495 islice_traverse(isliceobject *lz, visitproc visit, void *arg) 1496 { 1497 Py_VISIT(lz->it); 1498 return 0; 1499 } 1500 1501 static PyObject * 1502 islice_next(isliceobject *lz) 1503 { 1504 PyObject *item; 1505 PyObject *it = lz->it; 1506 Py_ssize_t stop = lz->stop; 1507 Py_ssize_t oldnext; 1508 PyObject *(*iternext)(PyObject *); 1509 1510 if (it == NULL) 1511 return NULL; 1512 1513 iternext = *Py_TYPE(it)->tp_iternext; 1514 while (lz->cnt < lz->next) { 1515 item = iternext(it); 1516 if (item == NULL) 1517 goto empty; 1518 Py_DECREF(item); 1519 lz->cnt++; 1520 } 1521 if (stop != -1 && lz->cnt >= stop) 1522 goto empty; 1523 item = iternext(it); 1524 if (item == NULL) 1525 goto empty; 1526 lz->cnt++; 1527 oldnext = lz->next; 1528 /* The (size_t) cast below avoids the danger of undefined 1529 behaviour from signed integer overflow. */ 1530 lz->next += (size_t)lz->step; 1531 if (lz->next < oldnext || (stop != -1 && lz->next > stop)) 1532 lz->next = stop; 1533 return item; 1534 1535 empty: 1536 Py_CLEAR(lz->it); 1537 return NULL; 1538 } 1539 1540 static PyObject * 1541 islice_reduce(isliceobject *lz) 1542 { 1543 /* When unpickled, generate a new object with the same bounds, 1544 * then 'setstate' with the next and count 1545 */ 1546 PyObject *stop; 1547 1548 if (lz->it == NULL) { 1549 PyObject *empty_list; 1550 PyObject *empty_it; 1551 empty_list = PyList_New(0); 1552 if (empty_list == NULL) 1553 return NULL; 1554 empty_it = PyObject_GetIter(empty_list); 1555 Py_DECREF(empty_list); 1556 if (empty_it == NULL) 1557 return NULL; 1558 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0); 1559 } 1560 if (lz->stop == -1) { 1561 stop = Py_None; 1562 Py_INCREF(stop); 1563 } else { 1564 stop = PyLong_FromSsize_t(lz->stop); 1565 if (stop == NULL) 1566 return NULL; 1567 } 1568 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz), 1569 lz->it, lz->next, stop, lz->step, 1570 lz->cnt); 1571 } 1572 1573 static PyObject * 1574 islice_setstate(isliceobject *lz, PyObject *state) 1575 { 1576 Py_ssize_t cnt = PyLong_AsSsize_t(state); 1577 1578 if (cnt == -1 && PyErr_Occurred()) 1579 return NULL; 1580 lz->cnt = cnt; 1581 Py_RETURN_NONE; 1582 } 1583 1584 static PyMethodDef islice_methods[] = { 1585 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS, 1586 reduce_doc}, 1587 {"__setstate__", (PyCFunction)islice_setstate, METH_O, 1588 setstate_doc}, 1589 {NULL, NULL} /* sentinel */ 1590 }; 1591 1592 PyDoc_STRVAR(islice_doc, 1593 "islice(iterable, stop) --> islice object\n\ 1594 islice(iterable, start, stop[, step]) --> islice object\n\ 1595 \n\ 1596 Return an iterator whose next() method returns selected values from an\n\ 1597 iterable. If start is specified, will skip all preceding elements;\n\ 1598 otherwise, start defaults to zero. Step defaults to one. If\n\ 1599 specified as another value, step determines how many values are \n\ 1600 skipped between successive calls. Works like a slice() on a list\n\ 1601 but returns an iterator."); 1602 1603 static PyTypeObject islice_type = { 1604 PyVarObject_HEAD_INIT(NULL, 0) 1605 "itertools.islice", /* tp_name */ 1606 sizeof(isliceobject), /* tp_basicsize */ 1607 0, /* tp_itemsize */ 1608 /* methods */ 1609 (destructor)islice_dealloc, /* tp_dealloc */ 1610 0, /* tp_print */ 1611 0, /* tp_getattr */ 1612 0, /* tp_setattr */ 1613 0, /* tp_reserved */ 1614 0, /* tp_repr */ 1615 0, /* tp_as_number */ 1616 0, /* tp_as_sequence */ 1617 0, /* tp_as_mapping */ 1618 0, /* tp_hash */ 1619 0, /* tp_call */ 1620 0, /* tp_str */ 1621 PyObject_GenericGetAttr, /* tp_getattro */ 1622 0, /* tp_setattro */ 1623 0, /* tp_as_buffer */ 1624 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 1625 Py_TPFLAGS_BASETYPE, /* tp_flags */ 1626 islice_doc, /* tp_doc */ 1627 (traverseproc)islice_traverse, /* tp_traverse */ 1628 0, /* tp_clear */ 1629 0, /* tp_richcompare */ 1630 0, /* tp_weaklistoffset */ 1631 PyObject_SelfIter, /* tp_iter */ 1632 (iternextfunc)islice_next, /* tp_iternext */ 1633 islice_methods, /* tp_methods */ 1634 0, /* tp_members */ 1635 0, /* tp_getset */ 1636 0, /* tp_base */ 1637 0, /* tp_dict */ 1638 0, /* tp_descr_get */ 1639 0, /* tp_descr_set */ 1640 0, /* tp_dictoffset */ 1641 0, /* tp_init */ 1642 0, /* tp_alloc */ 1643 islice_new, /* tp_new */ 1644 PyObject_GC_Del, /* tp_free */ 1645 }; 1646 1647 1648 /* starmap object ************************************************************/ 1649 1650 typedef struct { 1651 PyObject_HEAD 1652 PyObject *func; 1653 PyObject *it; 1654 } starmapobject; 1655 1656 static PyTypeObject starmap_type; 1657 1658 static PyObject * 1659 starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 1660 { 1661 PyObject *func, *seq; 1662 PyObject *it; 1663 starmapobject *lz; 1664 1665 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds)) 1666 return NULL; 1667 1668 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq)) 1669 return NULL; 1670 1671 /* Get iterator. */ 1672 it = PyObject_GetIter(seq); 1673 if (it == NULL) 1674 return NULL; 1675 1676 /* create starmapobject structure */ 1677 lz = (starmapobject *)type->tp_alloc(type, 0); 1678 if (lz == NULL) { 1679 Py_DECREF(it); 1680 return NULL; 1681 } 1682 Py_INCREF(func); 1683 lz->func = func; 1684 lz->it = it; 1685 1686 return (PyObject *)lz; 1687 } 1688 1689 static void 1690 starmap_dealloc(starmapobject *lz) 1691 { 1692 PyObject_GC_UnTrack(lz); 1693 Py_XDECREF(lz->func); 1694 Py_XDECREF(lz->it); 1695 Py_TYPE(lz)->tp_free(lz); 1696 } 1697 1698 static int 1699 starmap_traverse(starmapobject *lz, visitproc visit, void *arg) 1700 { 1701 Py_VISIT(lz->it); 1702 Py_VISIT(lz->func); 1703 return 0; 1704 } 1705 1706 static PyObject * 1707 starmap_next(starmapobject *lz) 1708 { 1709 PyObject *args; 1710 PyObject *result; 1711 PyObject *it = lz->it; 1712 1713 args = (*Py_TYPE(it)->tp_iternext)(it); 1714 if (args == NULL) 1715 return NULL; 1716 if (!PyTuple_CheckExact(args)) { 1717 PyObject *newargs = PySequence_Tuple(args); 1718 Py_DECREF(args); 1719 if (newargs == NULL) 1720 return NULL; 1721 args = newargs; 1722 } 1723 result = PyObject_Call(lz->func, args, NULL); 1724 Py_DECREF(args); 1725 return result; 1726 } 1727 1728 static PyObject * 1729 starmap_reduce(starmapobject *lz) 1730 { 1731 /* Just pickle the iterator */ 1732 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it); 1733 } 1734 1735 static PyMethodDef starmap_methods[] = { 1736 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS, 1737 reduce_doc}, 1738 {NULL, NULL} /* sentinel */ 1739 }; 1740 1741 PyDoc_STRVAR(starmap_doc, 1742 "starmap(function, sequence) --> starmap object\n\ 1743 \n\ 1744 Return an iterator whose values are returned from the function evaluated\n\ 1745 with an argument tuple taken from the given sequence."); 1746 1747 static PyTypeObject starmap_type = { 1748 PyVarObject_HEAD_INIT(NULL, 0) 1749 "itertools.starmap", /* tp_name */ 1750 sizeof(starmapobject), /* tp_basicsize */ 1751 0, /* tp_itemsize */ 1752 /* methods */ 1753 (destructor)starmap_dealloc, /* tp_dealloc */ 1754 0, /* tp_print */ 1755 0, /* tp_getattr */ 1756 0, /* tp_setattr */ 1757 0, /* tp_reserved */ 1758 0, /* tp_repr */ 1759 0, /* tp_as_number */ 1760 0, /* tp_as_sequence */ 1761 0, /* tp_as_mapping */ 1762 0, /* tp_hash */ 1763 0, /* tp_call */ 1764 0, /* tp_str */ 1765 PyObject_GenericGetAttr, /* tp_getattro */ 1766 0, /* tp_setattro */ 1767 0, /* tp_as_buffer */ 1768 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 1769 Py_TPFLAGS_BASETYPE, /* tp_flags */ 1770 starmap_doc, /* tp_doc */ 1771 (traverseproc)starmap_traverse, /* tp_traverse */ 1772 0, /* tp_clear */ 1773 0, /* tp_richcompare */ 1774 0, /* tp_weaklistoffset */ 1775 PyObject_SelfIter, /* tp_iter */ 1776 (iternextfunc)starmap_next, /* tp_iternext */ 1777 starmap_methods, /* tp_methods */ 1778 0, /* tp_members */ 1779 0, /* tp_getset */ 1780 0, /* tp_base */ 1781 0, /* tp_dict */ 1782 0, /* tp_descr_get */ 1783 0, /* tp_descr_set */ 1784 0, /* tp_dictoffset */ 1785 0, /* tp_init */ 1786 0, /* tp_alloc */ 1787 starmap_new, /* tp_new */ 1788 PyObject_GC_Del, /* tp_free */ 1789 }; 1790 1791 1792 /* chain object **************************************************************/ 1793 1794 typedef struct { 1795 PyObject_HEAD 1796 PyObject *source; /* Iterator over input iterables */ 1797 PyObject *active; /* Currently running input iterator */ 1798 } chainobject; 1799 1800 static PyTypeObject chain_type; 1801 1802 static PyObject * 1803 chain_new_internal(PyTypeObject *type, PyObject *source) 1804 { 1805 chainobject *lz; 1806 1807 lz = (chainobject *)type->tp_alloc(type, 0); 1808 if (lz == NULL) { 1809 Py_DECREF(source); 1810 return NULL; 1811 } 1812 1813 lz->source = source; 1814 lz->active = NULL; 1815 return (PyObject *)lz; 1816 } 1817 1818 static PyObject * 1819 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 1820 { 1821 PyObject *source; 1822 1823 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds)) 1824 return NULL; 1825 1826 source = PyObject_GetIter(args); 1827 if (source == NULL) 1828 return NULL; 1829 1830 return chain_new_internal(type, source); 1831 } 1832 1833 static PyObject * 1834 chain_new_from_iterable(PyTypeObject *type, PyObject *arg) 1835 { 1836 PyObject *source; 1837 1838 source = PyObject_GetIter(arg); 1839 if (source == NULL) 1840 return NULL; 1841 1842 return chain_new_internal(type, source); 1843 } 1844 1845 static void 1846 chain_dealloc(chainobject *lz) 1847 { 1848 PyObject_GC_UnTrack(lz); 1849 Py_XDECREF(lz->active); 1850 Py_XDECREF(lz->source); 1851 Py_TYPE(lz)->tp_free(lz); 1852 } 1853 1854 static int 1855 chain_traverse(chainobject *lz, visitproc visit, void *arg) 1856 { 1857 Py_VISIT(lz->source); 1858 Py_VISIT(lz->active); 1859 return 0; 1860 } 1861 1862 static PyObject * 1863 chain_next(chainobject *lz) 1864 { 1865 PyObject *item; 1866 1867 if (lz->source == NULL) 1868 return NULL; /* already stopped */ 1869 1870 if (lz->active == NULL) { 1871 PyObject *iterable = PyIter_Next(lz->source); 1872 if (iterable == NULL) { 1873 Py_CLEAR(lz->source); 1874 return NULL; /* no more input sources */ 1875 } 1876 lz->active = PyObject_GetIter(iterable); 1877 Py_DECREF(iterable); 1878 if (lz->active == NULL) { 1879 Py_CLEAR(lz->source); 1880 return NULL; /* input not iterable */ 1881 } 1882 } 1883 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active); 1884 if (item != NULL) 1885 return item; 1886 if (PyErr_Occurred()) { 1887 if (PyErr_ExceptionMatches(PyExc_StopIteration)) 1888 PyErr_Clear(); 1889 else 1890 return NULL; /* input raised an exception */ 1891 } 1892 Py_CLEAR(lz->active); 1893 return chain_next(lz); /* recurse and use next active */ 1894 } 1895 1896 static PyObject * 1897 chain_reduce(chainobject *lz) 1898 { 1899 if (lz->source) { 1900 /* we can't pickle function objects (itertools.from_iterable) so 1901 * we must use setstate to replace the iterable. One day we 1902 * will fix pickling of functions 1903 */ 1904 if (lz->active) { 1905 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active); 1906 } else { 1907 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source); 1908 } 1909 } else { 1910 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */ 1911 } 1912 return NULL; 1913 } 1914 1915 static PyObject * 1916 chain_setstate(chainobject *lz, PyObject *state) 1917 { 1918 PyObject *source, *active=NULL; 1919 1920 if (!PyTuple_Check(state)) { 1921 PyErr_SetString(PyExc_TypeError, "state is not a tuple"); 1922 return NULL; 1923 } 1924 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) { 1925 return NULL; 1926 } 1927 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) { 1928 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators."); 1929 return NULL; 1930 } 1931 1932 Py_INCREF(source); 1933 Py_XSETREF(lz->source, source); 1934 Py_XINCREF(active); 1935 Py_XSETREF(lz->active, active); 1936 Py_RETURN_NONE; 1937 } 1938 1939 PyDoc_STRVAR(chain_doc, 1940 "chain(*iterables) --> chain object\n\ 1941 \n\ 1942 Return a chain object whose .__next__() method returns elements from the\n\ 1943 first iterable until it is exhausted, then elements from the next\n\ 1944 iterable, until all of the iterables are exhausted."); 1945 1946 PyDoc_STRVAR(chain_from_iterable_doc, 1947 "chain.from_iterable(iterable) --> chain object\n\ 1948 \n\ 1949 Alternate chain() constructor taking a single iterable argument\n\ 1950 that evaluates lazily."); 1951 1952 static PyMethodDef chain_methods[] = { 1953 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS, 1954 chain_from_iterable_doc}, 1955 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS, 1956 reduce_doc}, 1957 {"__setstate__", (PyCFunction)chain_setstate, METH_O, 1958 setstate_doc}, 1959 {NULL, NULL} /* sentinel */ 1960 }; 1961 1962 static PyTypeObject chain_type = { 1963 PyVarObject_HEAD_INIT(NULL, 0) 1964 "itertools.chain", /* tp_name */ 1965 sizeof(chainobject), /* tp_basicsize */ 1966 0, /* tp_itemsize */ 1967 /* methods */ 1968 (destructor)chain_dealloc, /* tp_dealloc */ 1969 0, /* tp_print */ 1970 0, /* tp_getattr */ 1971 0, /* tp_setattr */ 1972 0, /* tp_reserved */ 1973 0, /* tp_repr */ 1974 0, /* tp_as_number */ 1975 0, /* tp_as_sequence */ 1976 0, /* tp_as_mapping */ 1977 0, /* tp_hash */ 1978 0, /* tp_call */ 1979 0, /* tp_str */ 1980 PyObject_GenericGetAttr, /* tp_getattro */ 1981 0, /* tp_setattro */ 1982 0, /* tp_as_buffer */ 1983 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 1984 Py_TPFLAGS_BASETYPE, /* tp_flags */ 1985 chain_doc, /* tp_doc */ 1986 (traverseproc)chain_traverse, /* tp_traverse */ 1987 0, /* tp_clear */ 1988 0, /* tp_richcompare */ 1989 0, /* tp_weaklistoffset */ 1990 PyObject_SelfIter, /* tp_iter */ 1991 (iternextfunc)chain_next, /* tp_iternext */ 1992 chain_methods, /* tp_methods */ 1993 0, /* tp_members */ 1994 0, /* tp_getset */ 1995 0, /* tp_base */ 1996 0, /* tp_dict */ 1997 0, /* tp_descr_get */ 1998 0, /* tp_descr_set */ 1999 0, /* tp_dictoffset */ 2000 0, /* tp_init */ 2001 0, /* tp_alloc */ 2002 chain_new, /* tp_new */ 2003 PyObject_GC_Del, /* tp_free */ 2004 }; 2005 2006 2007 /* product object ************************************************************/ 2008 2009 typedef struct { 2010 PyObject_HEAD 2011 PyObject *pools; /* tuple of pool tuples */ 2012 Py_ssize_t *indices; /* one index per pool */ 2013 PyObject *result; /* most recently returned result tuple */ 2014 int stopped; /* set to 1 when the iterator is exhausted */ 2015 } productobject; 2016 2017 static PyTypeObject product_type; 2018 2019 static PyObject * 2020 product_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 2021 { 2022 productobject *lz; 2023 Py_ssize_t nargs, npools, repeat=1; 2024 PyObject *pools = NULL; 2025 Py_ssize_t *indices = NULL; 2026 Py_ssize_t i; 2027 2028 if (kwds != NULL) { 2029 char *kwlist[] = {"repeat", 0}; 2030 PyObject *tmpargs = PyTuple_New(0); 2031 if (tmpargs == NULL) 2032 return NULL; 2033 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", 2034 kwlist, &repeat)) { 2035 Py_DECREF(tmpargs); 2036 return NULL; 2037 } 2038 Py_DECREF(tmpargs); 2039 if (repeat < 0) { 2040 PyErr_SetString(PyExc_ValueError, 2041 "repeat argument cannot be negative"); 2042 return NULL; 2043 } 2044 } 2045 2046 assert(PyTuple_CheckExact(args)); 2047 if (repeat == 0) { 2048 nargs = 0; 2049 } else { 2050 nargs = PyTuple_GET_SIZE(args); 2051 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) { 2052 PyErr_SetString(PyExc_OverflowError, "repeat argument too large"); 2053 return NULL; 2054 } 2055 } 2056 npools = nargs * repeat; 2057 2058 indices = PyMem_New(Py_ssize_t, npools); 2059 if (indices == NULL) { 2060 PyErr_NoMemory(); 2061 goto error; 2062 } 2063 2064 pools = PyTuple_New(npools); 2065 if (pools == NULL) 2066 goto error; 2067 2068 for (i=0; i < nargs ; ++i) { 2069 PyObject *item = PyTuple_GET_ITEM(args, i); 2070 PyObject *pool = PySequence_Tuple(item); 2071 if (pool == NULL) 2072 goto error; 2073 PyTuple_SET_ITEM(pools, i, pool); 2074 indices[i] = 0; 2075 } 2076 for ( ; i < npools; ++i) { 2077 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs); 2078 Py_INCREF(pool); 2079 PyTuple_SET_ITEM(pools, i, pool); 2080 indices[i] = 0; 2081 } 2082 2083 /* create productobject structure */ 2084 lz = (productobject *)type->tp_alloc(type, 0); 2085 if (lz == NULL) 2086 goto error; 2087 2088 lz->pools = pools; 2089 lz->indices = indices; 2090 lz->result = NULL; 2091 lz->stopped = 0; 2092 2093 return (PyObject *)lz; 2094 2095 error: 2096 if (indices != NULL) 2097 PyMem_Free(indices); 2098 Py_XDECREF(pools); 2099 return NULL; 2100 } 2101 2102 static void 2103 product_dealloc(productobject *lz) 2104 { 2105 PyObject_GC_UnTrack(lz); 2106 Py_XDECREF(lz->pools); 2107 Py_XDECREF(lz->result); 2108 if (lz->indices != NULL) 2109 PyMem_Free(lz->indices); 2110 Py_TYPE(lz)->tp_free(lz); 2111 } 2112 2113 static PyObject * 2114 product_sizeof(productobject *lz, void *unused) 2115 { 2116 Py_ssize_t res; 2117 2118 res = _PyObject_SIZE(Py_TYPE(lz)); 2119 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t); 2120 return PyLong_FromSsize_t(res); 2121 } 2122 2123 PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes."); 2124 2125 static int 2126 product_traverse(productobject *lz, visitproc visit, void *arg) 2127 { 2128 Py_VISIT(lz->pools); 2129 Py_VISIT(lz->result); 2130 return 0; 2131 } 2132 2133 static PyObject * 2134 product_next(productobject *lz) 2135 { 2136 PyObject *pool; 2137 PyObject *elem; 2138 PyObject *oldelem; 2139 PyObject *pools = lz->pools; 2140 PyObject *result = lz->result; 2141 Py_ssize_t npools = PyTuple_GET_SIZE(pools); 2142 Py_ssize_t i; 2143 2144 if (lz->stopped) 2145 return NULL; 2146 2147 if (result == NULL) { 2148 /* On the first pass, return an initial tuple filled with the 2149 first element from each pool. */ 2150 result = PyTuple_New(npools); 2151 if (result == NULL) 2152 goto empty; 2153 lz->result = result; 2154 for (i=0; i < npools; i++) { 2155 pool = PyTuple_GET_ITEM(pools, i); 2156 if (PyTuple_GET_SIZE(pool) == 0) 2157 goto empty; 2158 elem = PyTuple_GET_ITEM(pool, 0); 2159 Py_INCREF(elem); 2160 PyTuple_SET_ITEM(result, i, elem); 2161 } 2162 } else { 2163 Py_ssize_t *indices = lz->indices; 2164 2165 /* Copy the previous result tuple or re-use it if available */ 2166 if (Py_REFCNT(result) > 1) { 2167 PyObject *old_result = result; 2168 result = PyTuple_New(npools); 2169 if (result == NULL) 2170 goto empty; 2171 lz->result = result; 2172 for (i=0; i < npools; i++) { 2173 elem = PyTuple_GET_ITEM(old_result, i); 2174 Py_INCREF(elem); 2175 PyTuple_SET_ITEM(result, i, elem); 2176 } 2177 Py_DECREF(old_result); 2178 } 2179 /* Now, we've got the only copy so we can update it in-place */ 2180 assert (npools==0 || Py_REFCNT(result) == 1); 2181 2182 /* Update the pool indices right-to-left. Only advance to the 2183 next pool when the previous one rolls-over */ 2184 for (i=npools-1 ; i >= 0 ; i--) { 2185 pool = PyTuple_GET_ITEM(pools, i); 2186 indices[i]++; 2187 if (indices[i] == PyTuple_GET_SIZE(pool)) { 2188 /* Roll-over and advance to next pool */ 2189 indices[i] = 0; 2190 elem = PyTuple_GET_ITEM(pool, 0); 2191 Py_INCREF(elem); 2192 oldelem = PyTuple_GET_ITEM(result, i); 2193 PyTuple_SET_ITEM(result, i, elem); 2194 Py_DECREF(oldelem); 2195 } else { 2196 /* No rollover. Just increment and stop here. */ 2197 elem = PyTuple_GET_ITEM(pool, indices[i]); 2198 Py_INCREF(elem); 2199 oldelem = PyTuple_GET_ITEM(result, i); 2200 PyTuple_SET_ITEM(result, i, elem); 2201 Py_DECREF(oldelem); 2202 break; 2203 } 2204 } 2205 2206 /* If i is negative, then the indices have all rolled-over 2207 and we're done. */ 2208 if (i < 0) 2209 goto empty; 2210 } 2211 2212 Py_INCREF(result); 2213 return result; 2214 2215 empty: 2216 lz->stopped = 1; 2217 return NULL; 2218 } 2219 2220 static PyObject * 2221 product_reduce(productobject *lz) 2222 { 2223 if (lz->stopped) { 2224 return Py_BuildValue("O(())", Py_TYPE(lz)); 2225 } else if (lz->result == NULL) { 2226 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools); 2227 } else { 2228 PyObject *indices; 2229 Py_ssize_t n, i; 2230 2231 /* we must pickle the indices use them for setstate, and 2232 * additionally indicate that the iterator has started 2233 */ 2234 n = PyTuple_GET_SIZE(lz->pools); 2235 indices = PyTuple_New(n); 2236 if (indices == NULL) 2237 return NULL; 2238 for (i=0; i<n; i++){ 2239 PyObject* index = PyLong_FromSsize_t(lz->indices[i]); 2240 if (!index) { 2241 Py_DECREF(indices); 2242 return NULL; 2243 } 2244 PyTuple_SET_ITEM(indices, i, index); 2245 } 2246 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices); 2247 } 2248 } 2249 2250 static PyObject * 2251 product_setstate(productobject *lz, PyObject *state) 2252 { 2253 PyObject *result; 2254 Py_ssize_t n, i; 2255 2256 n = PyTuple_GET_SIZE(lz->pools); 2257 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) { 2258 PyErr_SetString(PyExc_ValueError, "invalid arguments"); 2259 return NULL; 2260 } 2261 for (i=0; i<n; i++) 2262 { 2263 PyObject* indexObject = PyTuple_GET_ITEM(state, i); 2264 Py_ssize_t index = PyLong_AsSsize_t(indexObject); 2265 PyObject* pool; 2266 Py_ssize_t poolsize; 2267 if (index < 0 && PyErr_Occurred()) 2268 return NULL; /* not an integer */ 2269 pool = PyTuple_GET_ITEM(lz->pools, i); 2270 poolsize = PyTuple_GET_SIZE(pool); 2271 if (poolsize == 0) { 2272 lz->stopped = 1; 2273 Py_RETURN_NONE; 2274 } 2275 /* clamp the index */ 2276 if (index < 0) 2277 index = 0; 2278 else if (index > poolsize-1) 2279 index = poolsize-1; 2280 lz->indices[i] = index; 2281 } 2282 2283 result = PyTuple_New(n); 2284 if (!result) 2285 return NULL; 2286 for (i=0; i<n; i++) { 2287 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i); 2288 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]); 2289 Py_INCREF(element); 2290 PyTuple_SET_ITEM(result, i, element); 2291 } 2292 Py_XSETREF(lz->result, result); 2293 Py_RETURN_NONE; 2294 } 2295 2296 static PyMethodDef product_methods[] = { 2297 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS, 2298 reduce_doc}, 2299 {"__setstate__", (PyCFunction)product_setstate, METH_O, 2300 setstate_doc}, 2301 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS, 2302 sizeof_doc}, 2303 {NULL, NULL} /* sentinel */ 2304 }; 2305 2306 PyDoc_STRVAR(product_doc, 2307 "product(*iterables, repeat=1) --> product object\n\ 2308 \n\ 2309 Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\ 2310 For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\ 2311 The leftmost iterators are in the outermost for-loop, so the output tuples\n\ 2312 cycle in a manner similar to an odometer (with the rightmost element changing\n\ 2313 on every iteration).\n\n\ 2314 To compute the product of an iterable with itself, specify the number\n\ 2315 of repetitions with the optional repeat keyword argument. For example,\n\ 2316 product(A, repeat=4) means the same as product(A, A, A, A).\n\n\ 2317 product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\ 2318 product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ..."); 2319 2320 static PyTypeObject product_type = { 2321 PyVarObject_HEAD_INIT(NULL, 0) 2322 "itertools.product", /* tp_name */ 2323 sizeof(productobject), /* tp_basicsize */ 2324 0, /* tp_itemsize */ 2325 /* methods */ 2326 (destructor)product_dealloc, /* tp_dealloc */ 2327 0, /* tp_print */ 2328 0, /* tp_getattr */ 2329 0, /* tp_setattr */ 2330 0, /* tp_reserved */ 2331 0, /* tp_repr */ 2332 0, /* tp_as_number */ 2333 0, /* tp_as_sequence */ 2334 0, /* tp_as_mapping */ 2335 0, /* tp_hash */ 2336 0, /* tp_call */ 2337 0, /* tp_str */ 2338 PyObject_GenericGetAttr, /* tp_getattro */ 2339 0, /* tp_setattro */ 2340 0, /* tp_as_buffer */ 2341 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 2342 Py_TPFLAGS_BASETYPE, /* tp_flags */ 2343 product_doc, /* tp_doc */ 2344 (traverseproc)product_traverse, /* tp_traverse */ 2345 0, /* tp_clear */ 2346 0, /* tp_richcompare */ 2347 0, /* tp_weaklistoffset */ 2348 PyObject_SelfIter, /* tp_iter */ 2349 (iternextfunc)product_next, /* tp_iternext */ 2350 product_methods, /* tp_methods */ 2351 0, /* tp_members */ 2352 0, /* tp_getset */ 2353 0, /* tp_base */ 2354 0, /* tp_dict */ 2355 0, /* tp_descr_get */ 2356 0, /* tp_descr_set */ 2357 0, /* tp_dictoffset */ 2358 0, /* tp_init */ 2359 0, /* tp_alloc */ 2360 product_new, /* tp_new */ 2361 PyObject_GC_Del, /* tp_free */ 2362 }; 2363 2364 2365 /* combinations object *******************************************************/ 2366 2367 typedef struct { 2368 PyObject_HEAD 2369 PyObject *pool; /* input converted to a tuple */ 2370 Py_ssize_t *indices; /* one index per result element */ 2371 PyObject *result; /* most recently returned result tuple */ 2372 Py_ssize_t r; /* size of result tuple */ 2373 int stopped; /* set to 1 when the iterator is exhausted */ 2374 } combinationsobject; 2375 2376 static PyTypeObject combinations_type; 2377 2378 static PyObject * 2379 combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 2380 { 2381 combinationsobject *co; 2382 Py_ssize_t n; 2383 Py_ssize_t r; 2384 PyObject *pool = NULL; 2385 PyObject *iterable = NULL; 2386 Py_ssize_t *indices = NULL; 2387 Py_ssize_t i; 2388 static char *kwargs[] = {"iterable", "r", NULL}; 2389 2390 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs, 2391 &iterable, &r)) 2392 return NULL; 2393 2394 pool = PySequence_Tuple(iterable); 2395 if (pool == NULL) 2396 goto error; 2397 n = PyTuple_GET_SIZE(pool); 2398 if (r < 0) { 2399 PyErr_SetString(PyExc_ValueError, "r must be non-negative"); 2400 goto error; 2401 } 2402 2403 indices = PyMem_New(Py_ssize_t, r); 2404 if (indices == NULL) { 2405 PyErr_NoMemory(); 2406 goto error; 2407 } 2408 2409 for (i=0 ; i<r ; i++) 2410 indices[i] = i; 2411 2412 /* create combinationsobject structure */ 2413 co = (combinationsobject *)type->tp_alloc(type, 0); 2414 if (co == NULL) 2415 goto error; 2416 2417 co->pool = pool; 2418 co->indices = indices; 2419 co->result = NULL; 2420 co->r = r; 2421 co->stopped = r > n ? 1 : 0; 2422 2423 return (PyObject *)co; 2424 2425 error: 2426 if (indices != NULL) 2427 PyMem_Free(indices); 2428 Py_XDECREF(pool); 2429 return NULL; 2430 } 2431 2432 static void 2433 combinations_dealloc(combinationsobject *co) 2434 { 2435 PyObject_GC_UnTrack(co); 2436 Py_XDECREF(co->pool); 2437 Py_XDECREF(co->result); 2438 if (co->indices != NULL) 2439 PyMem_Free(co->indices); 2440 Py_TYPE(co)->tp_free(co); 2441 } 2442 2443 static PyObject * 2444 combinations_sizeof(combinationsobject *co, void *unused) 2445 { 2446 Py_ssize_t res; 2447 2448 res = _PyObject_SIZE(Py_TYPE(co)); 2449 res += co->r * sizeof(Py_ssize_t); 2450 return PyLong_FromSsize_t(res); 2451 } 2452 2453 static int 2454 combinations_traverse(combinationsobject *co, visitproc visit, void *arg) 2455 { 2456 Py_VISIT(co->pool); 2457 Py_VISIT(co->result); 2458 return 0; 2459 } 2460 2461 static PyObject * 2462 combinations_next(combinationsobject *co) 2463 { 2464 PyObject *elem; 2465 PyObject *oldelem; 2466 PyObject *pool = co->pool; 2467 Py_ssize_t *indices = co->indices; 2468 PyObject *result = co->result; 2469 Py_ssize_t n = PyTuple_GET_SIZE(pool); 2470 Py_ssize_t r = co->r; 2471 Py_ssize_t i, j, index; 2472 2473 if (co->stopped) 2474 return NULL; 2475 2476 if (result == NULL) { 2477 /* On the first pass, initialize result tuple using the indices */ 2478 result = PyTuple_New(r); 2479 if (result == NULL) 2480 goto empty; 2481 co->result = result; 2482 for (i=0; i<r ; i++) { 2483 index = indices[i]; 2484 elem = PyTuple_GET_ITEM(pool, index); 2485 Py_INCREF(elem); 2486 PyTuple_SET_ITEM(result, i, elem); 2487 } 2488 } else { 2489 /* Copy the previous result tuple or re-use it if available */ 2490 if (Py_REFCNT(result) > 1) { 2491 PyObject *old_result = result; 2492 result = PyTuple_New(r); 2493 if (result == NULL) 2494 goto empty; 2495 co->result = result; 2496 for (i=0; i<r ; i++) { 2497 elem = PyTuple_GET_ITEM(old_result, i); 2498 Py_INCREF(elem); 2499 PyTuple_SET_ITEM(result, i, elem); 2500 } 2501 Py_DECREF(old_result); 2502 } 2503 /* Now, we've got the only copy so we can update it in-place 2504 * CPython's empty tuple is a singleton and cached in 2505 * PyTuple's freelist. 2506 */ 2507 assert(r == 0 || Py_REFCNT(result) == 1); 2508 2509 /* Scan indices right-to-left until finding one that is not 2510 at its maximum (i + n - r). */ 2511 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--) 2512 ; 2513 2514 /* If i is negative, then the indices are all at 2515 their maximum value and we're done. */ 2516 if (i < 0) 2517 goto empty; 2518 2519 /* Increment the current index which we know is not at its 2520 maximum. Then move back to the right setting each index 2521 to its lowest possible value (one higher than the index 2522 to its left -- this maintains the sort order invariant). */ 2523 indices[i]++; 2524 for (j=i+1 ; j<r ; j++) 2525 indices[j] = indices[j-1] + 1; 2526 2527 /* Update the result tuple for the new indices 2528 starting with i, the leftmost index that changed */ 2529 for ( ; i<r ; i++) { 2530 index = indices[i]; 2531 elem = PyTuple_GET_ITEM(pool, index); 2532 Py_INCREF(elem); 2533 oldelem = PyTuple_GET_ITEM(result, i); 2534 PyTuple_SET_ITEM(result, i, elem); 2535 Py_DECREF(oldelem); 2536 } 2537 } 2538 2539 Py_INCREF(result); 2540 return result; 2541 2542 empty: 2543 co->stopped = 1; 2544 return NULL; 2545 } 2546 2547 static PyObject * 2548 combinations_reduce(combinationsobject *lz) 2549 { 2550 if (lz->result == NULL) { 2551 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r); 2552 } else if (lz->stopped) { 2553 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r); 2554 } else { 2555 PyObject *indices; 2556 Py_ssize_t i; 2557 2558 /* we must pickle the indices and use them for setstate */ 2559 indices = PyTuple_New(lz->r); 2560 if (!indices) 2561 return NULL; 2562 for (i=0; i<lz->r; i++) 2563 { 2564 PyObject* index = PyLong_FromSsize_t(lz->indices[i]); 2565 if (!index) { 2566 Py_DECREF(indices); 2567 return NULL; 2568 } 2569 PyTuple_SET_ITEM(indices, i, index); 2570 } 2571 2572 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices); 2573 } 2574 } 2575 2576 static PyObject * 2577 combinations_setstate(combinationsobject *lz, PyObject *state) 2578 { 2579 PyObject *result; 2580 Py_ssize_t i; 2581 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool); 2582 2583 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) { 2584 PyErr_SetString(PyExc_ValueError, "invalid arguments"); 2585 return NULL; 2586 } 2587 2588 for (i=0; i<lz->r; i++) { 2589 Py_ssize_t max; 2590 PyObject* indexObject = PyTuple_GET_ITEM(state, i); 2591 Py_ssize_t index = PyLong_AsSsize_t(indexObject); 2592 2593 if (index == -1 && PyErr_Occurred()) 2594 return NULL; /* not an integer */ 2595 max = i + n - lz->r; 2596 /* clamp the index (beware of negative max) */ 2597 if (index > max) 2598 index = max; 2599 if (index < 0) 2600 index = 0; 2601 lz->indices[i] = index; 2602 } 2603 2604 result = PyTuple_New(lz->r); 2605 if (result == NULL) 2606 return NULL; 2607 for (i=0; i<lz->r; i++) { 2608 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]); 2609 Py_INCREF(element); 2610 PyTuple_SET_ITEM(result, i, element); 2611 } 2612 2613 Py_XSETREF(lz->result, result); 2614 Py_RETURN_NONE; 2615 } 2616 2617 static PyMethodDef combinations_methods[] = { 2618 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS, 2619 reduce_doc}, 2620 {"__setstate__", (PyCFunction)combinations_setstate, METH_O, 2621 setstate_doc}, 2622 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS, 2623 sizeof_doc}, 2624 {NULL, NULL} /* sentinel */ 2625 }; 2626 2627 PyDoc_STRVAR(combinations_doc, 2628 "combinations(iterable, r) --> combinations object\n\ 2629 \n\ 2630 Return successive r-length combinations of elements in the iterable.\n\n\ 2631 combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)"); 2632 2633 static PyTypeObject combinations_type = { 2634 PyVarObject_HEAD_INIT(NULL, 0) 2635 "itertools.combinations", /* tp_name */ 2636 sizeof(combinationsobject), /* tp_basicsize */ 2637 0, /* tp_itemsize */ 2638 /* methods */ 2639 (destructor)combinations_dealloc, /* tp_dealloc */ 2640 0, /* tp_print */ 2641 0, /* tp_getattr */ 2642 0, /* tp_setattr */ 2643 0, /* tp_reserved */ 2644 0, /* tp_repr */ 2645 0, /* tp_as_number */ 2646 0, /* tp_as_sequence */ 2647 0, /* tp_as_mapping */ 2648 0, /* tp_hash */ 2649 0, /* tp_call */ 2650 0, /* tp_str */ 2651 PyObject_GenericGetAttr, /* tp_getattro */ 2652 0, /* tp_setattro */ 2653 0, /* tp_as_buffer */ 2654 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 2655 Py_TPFLAGS_BASETYPE, /* tp_flags */ 2656 combinations_doc, /* tp_doc */ 2657 (traverseproc)combinations_traverse,/* tp_traverse */ 2658 0, /* tp_clear */ 2659 0, /* tp_richcompare */ 2660 0, /* tp_weaklistoffset */ 2661 PyObject_SelfIter, /* tp_iter */ 2662 (iternextfunc)combinations_next, /* tp_iternext */ 2663 combinations_methods, /* tp_methods */ 2664 0, /* tp_members */ 2665 0, /* tp_getset */ 2666 0, /* tp_base */ 2667 0, /* tp_dict */ 2668 0, /* tp_descr_get */ 2669 0, /* tp_descr_set */ 2670 0, /* tp_dictoffset */ 2671 0, /* tp_init */ 2672 0, /* tp_alloc */ 2673 combinations_new, /* tp_new */ 2674 PyObject_GC_Del, /* tp_free */ 2675 }; 2676 2677 2678 /* combinations with replacement object **************************************/ 2679 2680 /* Equivalent to: 2681 2682 def combinations_with_replacement(iterable, r): 2683 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC" 2684 # number items returned: (n+r-1)! / r! / (n-1)! 2685 pool = tuple(iterable) 2686 n = len(pool) 2687 indices = [0] * r 2688 yield tuple(pool[i] for i in indices) 2689 while 1: 2690 for i in reversed(range(r)): 2691 if indices[i] != n - 1: 2692 break 2693 else: 2694 return 2695 indices[i:] = [indices[i] + 1] * (r - i) 2696 yield tuple(pool[i] for i in indices) 2697 2698 def combinations_with_replacement2(iterable, r): 2699 'Alternate version that filters from product()' 2700 pool = tuple(iterable) 2701 n = len(pool) 2702 for indices in product(range(n), repeat=r): 2703 if sorted(indices) == list(indices): 2704 yield tuple(pool[i] for i in indices) 2705 */ 2706 typedef struct { 2707 PyObject_HEAD 2708 PyObject *pool; /* input converted to a tuple */ 2709 Py_ssize_t *indices; /* one index per result element */ 2710 PyObject *result; /* most recently returned result tuple */ 2711 Py_ssize_t r; /* size of result tuple */ 2712 int stopped; /* set to 1 when the cwr iterator is exhausted */ 2713 } cwrobject; 2714 2715 static PyTypeObject cwr_type; 2716 2717 static PyObject * 2718 cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 2719 { 2720 cwrobject *co; 2721 Py_ssize_t n; 2722 Py_ssize_t r; 2723 PyObject *pool = NULL; 2724 PyObject *iterable = NULL; 2725 Py_ssize_t *indices = NULL; 2726 Py_ssize_t i; 2727 static char *kwargs[] = {"iterable", "r", NULL}; 2728 2729 if (!PyArg_ParseTupleAndKeywords(args, kwds, 2730 "On:combinations_with_replacement", 2731 kwargs, &iterable, &r)) 2732 return NULL; 2733 2734 pool = PySequence_Tuple(iterable); 2735 if (pool == NULL) 2736 goto error; 2737 n = PyTuple_GET_SIZE(pool); 2738 if (r < 0) { 2739 PyErr_SetString(PyExc_ValueError, "r must be non-negative"); 2740 goto error; 2741 } 2742 2743 indices = PyMem_New(Py_ssize_t, r); 2744 if (indices == NULL) { 2745 PyErr_NoMemory(); 2746 goto error; 2747 } 2748 2749 for (i=0 ; i<r ; i++) 2750 indices[i] = 0; 2751 2752 /* create cwrobject structure */ 2753 co = (cwrobject *)type->tp_alloc(type, 0); 2754 if (co == NULL) 2755 goto error; 2756 2757 co->pool = pool; 2758 co->indices = indices; 2759 co->result = NULL; 2760 co->r = r; 2761 co->stopped = !n && r; 2762 2763 return (PyObject *)co; 2764 2765 error: 2766 if (indices != NULL) 2767 PyMem_Free(indices); 2768 Py_XDECREF(pool); 2769 return NULL; 2770 } 2771 2772 static void 2773 cwr_dealloc(cwrobject *co) 2774 { 2775 PyObject_GC_UnTrack(co); 2776 Py_XDECREF(co->pool); 2777 Py_XDECREF(co->result); 2778 if (co->indices != NULL) 2779 PyMem_Free(co->indices); 2780 Py_TYPE(co)->tp_free(co); 2781 } 2782 2783 static PyObject * 2784 cwr_sizeof(cwrobject *co, void *unused) 2785 { 2786 Py_ssize_t res; 2787 2788 res = _PyObject_SIZE(Py_TYPE(co)); 2789 res += co->r * sizeof(Py_ssize_t); 2790 return PyLong_FromSsize_t(res); 2791 } 2792 2793 static int 2794 cwr_traverse(cwrobject *co, visitproc visit, void *arg) 2795 { 2796 Py_VISIT(co->pool); 2797 Py_VISIT(co->result); 2798 return 0; 2799 } 2800 2801 static PyObject * 2802 cwr_next(cwrobject *co) 2803 { 2804 PyObject *elem; 2805 PyObject *oldelem; 2806 PyObject *pool = co->pool; 2807 Py_ssize_t *indices = co->indices; 2808 PyObject *result = co->result; 2809 Py_ssize_t n = PyTuple_GET_SIZE(pool); 2810 Py_ssize_t r = co->r; 2811 Py_ssize_t i, index; 2812 2813 if (co->stopped) 2814 return NULL; 2815 2816 if (result == NULL) { 2817 /* On the first pass, initialize result tuple with pool[0] */ 2818 result = PyTuple_New(r); 2819 if (result == NULL) 2820 goto empty; 2821 co->result = result; 2822 if (n > 0) { 2823 elem = PyTuple_GET_ITEM(pool, 0); 2824 for (i=0; i<r ; i++) { 2825 assert(indices[i] == 0); 2826 Py_INCREF(elem); 2827 PyTuple_SET_ITEM(result, i, elem); 2828 } 2829 } 2830 } else { 2831 /* Copy the previous result tuple or re-use it if available */ 2832 if (Py_REFCNT(result) > 1) { 2833 PyObject *old_result = result; 2834 result = PyTuple_New(r); 2835 if (result == NULL) 2836 goto empty; 2837 co->result = result; 2838 for (i=0; i<r ; i++) { 2839 elem = PyTuple_GET_ITEM(old_result, i); 2840 Py_INCREF(elem); 2841 PyTuple_SET_ITEM(result, i, elem); 2842 } 2843 Py_DECREF(old_result); 2844 } 2845 /* Now, we've got the only copy so we can update it in-place CPython's 2846 empty tuple is a singleton and cached in PyTuple's freelist. */ 2847 assert(r == 0 || Py_REFCNT(result) == 1); 2848 2849 /* Scan indices right-to-left until finding one that is not 2850 * at its maximum (n-1). */ 2851 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--) 2852 ; 2853 2854 /* If i is negative, then the indices are all at 2855 their maximum value and we're done. */ 2856 if (i < 0) 2857 goto empty; 2858 2859 /* Increment the current index which we know is not at its 2860 maximum. Then set all to the right to the same value. */ 2861 index = indices[i] + 1; 2862 assert(index < n); 2863 elem = PyTuple_GET_ITEM(pool, index); 2864 for ( ; i<r ; i++) { 2865 indices[i] = index; 2866 Py_INCREF(elem); 2867 oldelem = PyTuple_GET_ITEM(result, i); 2868 PyTuple_SET_ITEM(result, i, elem); 2869 Py_DECREF(oldelem); 2870 } 2871 } 2872 2873 Py_INCREF(result); 2874 return result; 2875 2876 empty: 2877 co->stopped = 1; 2878 return NULL; 2879 } 2880 2881 static PyObject * 2882 cwr_reduce(cwrobject *lz) 2883 { 2884 if (lz->result == NULL) { 2885 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r); 2886 } else if (lz->stopped) { 2887 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r); 2888 } else { 2889 PyObject *indices; 2890 Py_ssize_t i; 2891 2892 /* we must pickle the indices and use them for setstate */ 2893 indices = PyTuple_New(lz->r); 2894 if (!indices) 2895 return NULL; 2896 for (i=0; i<lz->r; i++) { 2897 PyObject* index = PyLong_FromSsize_t(lz->indices[i]); 2898 if (!index) { 2899 Py_DECREF(indices); 2900 return NULL; 2901 } 2902 PyTuple_SET_ITEM(indices, i, index); 2903 } 2904 2905 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices); 2906 } 2907 } 2908 2909 static PyObject * 2910 cwr_setstate(cwrobject *lz, PyObject *state) 2911 { 2912 PyObject *result; 2913 Py_ssize_t n, i; 2914 2915 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) 2916 { 2917 PyErr_SetString(PyExc_ValueError, "invalid arguments"); 2918 return NULL; 2919 } 2920 2921 n = PyTuple_GET_SIZE(lz->pool); 2922 for (i=0; i<lz->r; i++) { 2923 PyObject* indexObject = PyTuple_GET_ITEM(state, i); 2924 Py_ssize_t index = PyLong_AsSsize_t(indexObject); 2925 2926 if (index < 0 && PyErr_Occurred()) 2927 return NULL; /* not an integer */ 2928 /* clamp the index */ 2929 if (index < 0) 2930 index = 0; 2931 else if (index > n-1) 2932 index = n-1; 2933 lz->indices[i] = index; 2934 } 2935 result = PyTuple_New(lz->r); 2936 if (result == NULL) 2937 return NULL; 2938 for (i=0; i<lz->r; i++) { 2939 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]); 2940 Py_INCREF(element); 2941 PyTuple_SET_ITEM(result, i, element); 2942 } 2943 Py_XSETREF(lz->result, result); 2944 Py_RETURN_NONE; 2945 } 2946 2947 static PyMethodDef cwr_methods[] = { 2948 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS, 2949 reduce_doc}, 2950 {"__setstate__", (PyCFunction)cwr_setstate, METH_O, 2951 setstate_doc}, 2952 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS, 2953 sizeof_doc}, 2954 {NULL, NULL} /* sentinel */ 2955 }; 2956 2957 PyDoc_STRVAR(cwr_doc, 2958 "combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\ 2959 \n\ 2960 Return successive r-length combinations of elements in the iterable\n\ 2961 allowing individual elements to have successive repeats.\n\ 2962 combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"); 2963 2964 static PyTypeObject cwr_type = { 2965 PyVarObject_HEAD_INIT(NULL, 0) 2966 "itertools.combinations_with_replacement", /* tp_name */ 2967 sizeof(cwrobject), /* tp_basicsize */ 2968 0, /* tp_itemsize */ 2969 /* methods */ 2970 (destructor)cwr_dealloc, /* tp_dealloc */ 2971 0, /* tp_print */ 2972 0, /* tp_getattr */ 2973 0, /* tp_setattr */ 2974 0, /* tp_reserved */ 2975 0, /* tp_repr */ 2976 0, /* tp_as_number */ 2977 0, /* tp_as_sequence */ 2978 0, /* tp_as_mapping */ 2979 0, /* tp_hash */ 2980 0, /* tp_call */ 2981 0, /* tp_str */ 2982 PyObject_GenericGetAttr, /* tp_getattro */ 2983 0, /* tp_setattro */ 2984 0, /* tp_as_buffer */ 2985 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 2986 Py_TPFLAGS_BASETYPE, /* tp_flags */ 2987 cwr_doc, /* tp_doc */ 2988 (traverseproc)cwr_traverse, /* tp_traverse */ 2989 0, /* tp_clear */ 2990 0, /* tp_richcompare */ 2991 0, /* tp_weaklistoffset */ 2992 PyObject_SelfIter, /* tp_iter */ 2993 (iternextfunc)cwr_next, /* tp_iternext */ 2994 cwr_methods, /* tp_methods */ 2995 0, /* tp_members */ 2996 0, /* tp_getset */ 2997 0, /* tp_base */ 2998 0, /* tp_dict */ 2999 0, /* tp_descr_get */ 3000 0, /* tp_descr_set */ 3001 0, /* tp_dictoffset */ 3002 0, /* tp_init */ 3003 0, /* tp_alloc */ 3004 cwr_new, /* tp_new */ 3005 PyObject_GC_Del, /* tp_free */ 3006 }; 3007 3008 3009 /* permutations object ******************************************************** 3010 3011 def permutations(iterable, r=None): 3012 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)' 3013 pool = tuple(iterable) 3014 n = len(pool) 3015 r = n if r is None else r 3016 indices = range(n) 3017 cycles = range(n-r+1, n+1)[::-1] 3018 yield tuple(pool[i] for i in indices[:r]) 3019 while n: 3020 for i in reversed(range(r)): 3021 cycles[i] -= 1 3022 if cycles[i] == 0: 3023 indices[i:] = indices[i+1:] + indices[i:i+1] 3024 cycles[i] = n - i 3025 else: 3026 j = cycles[i] 3027 indices[i], indices[-j] = indices[-j], indices[i] 3028 yield tuple(pool[i] for i in indices[:r]) 3029 break 3030 else: 3031 return 3032 */ 3033 3034 typedef struct { 3035 PyObject_HEAD 3036 PyObject *pool; /* input converted to a tuple */ 3037 Py_ssize_t *indices; /* one index per element in the pool */ 3038 Py_ssize_t *cycles; /* one rollover counter per element in the result */ 3039 PyObject *result; /* most recently returned result tuple */ 3040 Py_ssize_t r; /* size of result tuple */ 3041 int stopped; /* set to 1 when the iterator is exhausted */ 3042 } permutationsobject; 3043 3044 static PyTypeObject permutations_type; 3045 3046 static PyObject * 3047 permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 3048 { 3049 permutationsobject *po; 3050 Py_ssize_t n; 3051 Py_ssize_t r; 3052 PyObject *robj = Py_None; 3053 PyObject *pool = NULL; 3054 PyObject *iterable = NULL; 3055 Py_ssize_t *indices = NULL; 3056 Py_ssize_t *cycles = NULL; 3057 Py_ssize_t i; 3058 static char *kwargs[] = {"iterable", "r", NULL}; 3059 3060 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs, 3061 &iterable, &robj)) 3062 return NULL; 3063 3064 pool = PySequence_Tuple(iterable); 3065 if (pool == NULL) 3066 goto error; 3067 n = PyTuple_GET_SIZE(pool); 3068 3069 r = n; 3070 if (robj != Py_None) { 3071 if (!PyLong_Check(robj)) { 3072 PyErr_SetString(PyExc_TypeError, "Expected int as r"); 3073 goto error; 3074 } 3075 r = PyLong_AsSsize_t(robj); 3076 if (r == -1 && PyErr_Occurred()) 3077 goto error; 3078 } 3079 if (r < 0) { 3080 PyErr_SetString(PyExc_ValueError, "r must be non-negative"); 3081 goto error; 3082 } 3083 3084 indices = PyMem_New(Py_ssize_t, n); 3085 cycles = PyMem_New(Py_ssize_t, r); 3086 if (indices == NULL || cycles == NULL) { 3087 PyErr_NoMemory(); 3088 goto error; 3089 } 3090 3091 for (i=0 ; i<n ; i++) 3092 indices[i] = i; 3093 for (i=0 ; i<r ; i++) 3094 cycles[i] = n - i; 3095 3096 /* create permutationsobject structure */ 3097 po = (permutationsobject *)type->tp_alloc(type, 0); 3098 if (po == NULL) 3099 goto error; 3100 3101 po->pool = pool; 3102 po->indices = indices; 3103 po->cycles = cycles; 3104 po->result = NULL; 3105 po->r = r; 3106 po->stopped = r > n ? 1 : 0; 3107 3108 return (PyObject *)po; 3109 3110 error: 3111 if (indices != NULL) 3112 PyMem_Free(indices); 3113 if (cycles != NULL) 3114 PyMem_Free(cycles); 3115 Py_XDECREF(pool); 3116 return NULL; 3117 } 3118 3119 static void 3120 permutations_dealloc(permutationsobject *po) 3121 { 3122 PyObject_GC_UnTrack(po); 3123 Py_XDECREF(po->pool); 3124 Py_XDECREF(po->result); 3125 PyMem_Free(po->indices); 3126 PyMem_Free(po->cycles); 3127 Py_TYPE(po)->tp_free(po); 3128 } 3129 3130 static PyObject * 3131 permutations_sizeof(permutationsobject *po, void *unused) 3132 { 3133 Py_ssize_t res; 3134 3135 res = _PyObject_SIZE(Py_TYPE(po)); 3136 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t); 3137 res += po->r * sizeof(Py_ssize_t); 3138 return PyLong_FromSsize_t(res); 3139 } 3140 3141 static int 3142 permutations_traverse(permutationsobject *po, visitproc visit, void *arg) 3143 { 3144 Py_VISIT(po->pool); 3145 Py_VISIT(po->result); 3146 return 0; 3147 } 3148 3149 static PyObject * 3150 permutations_next(permutationsobject *po) 3151 { 3152 PyObject *elem; 3153 PyObject *oldelem; 3154 PyObject *pool = po->pool; 3155 Py_ssize_t *indices = po->indices; 3156 Py_ssize_t *cycles = po->cycles; 3157 PyObject *result = po->result; 3158 Py_ssize_t n = PyTuple_GET_SIZE(pool); 3159 Py_ssize_t r = po->r; 3160 Py_ssize_t i, j, k, index; 3161 3162 if (po->stopped) 3163 return NULL; 3164 3165 if (result == NULL) { 3166 /* On the first pass, initialize result tuple using the indices */ 3167 result = PyTuple_New(r); 3168 if (result == NULL) 3169 goto empty; 3170 po->result = result; 3171 for (i=0; i<r ; i++) { 3172 index = indices[i]; 3173 elem = PyTuple_GET_ITEM(pool, index); 3174 Py_INCREF(elem); 3175 PyTuple_SET_ITEM(result, i, elem); 3176 } 3177 } else { 3178 if (n == 0) 3179 goto empty; 3180 3181 /* Copy the previous result tuple or re-use it if available */ 3182 if (Py_REFCNT(result) > 1) { 3183 PyObject *old_result = result; 3184 result = PyTuple_New(r); 3185 if (result == NULL) 3186 goto empty; 3187 po->result = result; 3188 for (i=0; i<r ; i++) { 3189 elem = PyTuple_GET_ITEM(old_result, i); 3190 Py_INCREF(elem); 3191 PyTuple_SET_ITEM(result, i, elem); 3192 } 3193 Py_DECREF(old_result); 3194 } 3195 /* Now, we've got the only copy so we can update it in-place */ 3196 assert(r == 0 || Py_REFCNT(result) == 1); 3197 3198 /* Decrement rightmost cycle, moving leftward upon zero rollover */ 3199 for (i=r-1 ; i>=0 ; i--) { 3200 cycles[i] -= 1; 3201 if (cycles[i] == 0) { 3202 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */ 3203 index = indices[i]; 3204 for (j=i ; j<n-1 ; j++) 3205 indices[j] = indices[j+1]; 3206 indices[n-1] = index; 3207 cycles[i] = n - i; 3208 } else { 3209 j = cycles[i]; 3210 index = indices[i]; 3211 indices[i] = indices[n-j]; 3212 indices[n-j] = index; 3213 3214 for (k=i; k<r ; k++) { 3215 /* start with i, the leftmost element that changed */ 3216 /* yield tuple(pool[k] for k in indices[:r]) */ 3217 index = indices[k]; 3218 elem = PyTuple_GET_ITEM(pool, index); 3219 Py_INCREF(elem); 3220 oldelem = PyTuple_GET_ITEM(result, k); 3221 PyTuple_SET_ITEM(result, k, elem); 3222 Py_DECREF(oldelem); 3223 } 3224 break; 3225 } 3226 } 3227 /* If i is negative, then the cycles have all 3228 rolled-over and we're done. */ 3229 if (i < 0) 3230 goto empty; 3231 } 3232 Py_INCREF(result); 3233 return result; 3234 3235 empty: 3236 po->stopped = 1; 3237 return NULL; 3238 } 3239 3240 static PyObject * 3241 permutations_reduce(permutationsobject *po) 3242 { 3243 if (po->result == NULL) { 3244 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r); 3245 } else if (po->stopped) { 3246 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r); 3247 } else { 3248 PyObject *indices=NULL, *cycles=NULL; 3249 Py_ssize_t n, i; 3250 3251 /* we must pickle the indices and cycles and use them for setstate */ 3252 n = PyTuple_GET_SIZE(po->pool); 3253 indices = PyTuple_New(n); 3254 if (indices == NULL) 3255 goto err; 3256 for (i=0; i<n; i++) { 3257 PyObject* index = PyLong_FromSsize_t(po->indices[i]); 3258 if (!index) 3259 goto err; 3260 PyTuple_SET_ITEM(indices, i, index); 3261 } 3262 3263 cycles = PyTuple_New(po->r); 3264 if (cycles == NULL) 3265 goto err; 3266 for (i=0 ; i<po->r ; i++) { 3267 PyObject* index = PyLong_FromSsize_t(po->cycles[i]); 3268 if (!index) 3269 goto err; 3270 PyTuple_SET_ITEM(cycles, i, index); 3271 } 3272 return Py_BuildValue("O(On)(NN)", Py_TYPE(po), 3273 po->pool, po->r, 3274 indices, cycles); 3275 err: 3276 Py_XDECREF(indices); 3277 Py_XDECREF(cycles); 3278 return NULL; 3279 } 3280 } 3281 3282 static PyObject * 3283 permutations_setstate(permutationsobject *po, PyObject *state) 3284 { 3285 PyObject *indices, *cycles, *result; 3286 Py_ssize_t n, i; 3287 3288 if (!PyTuple_Check(state)) { 3289 PyErr_SetString(PyExc_TypeError, "state is not a tuple"); 3290 return NULL; 3291 } 3292 if (!PyArg_ParseTuple(state, "O!O!", 3293 &PyTuple_Type, &indices, 3294 &PyTuple_Type, &cycles)) { 3295 return NULL; 3296 } 3297 3298 n = PyTuple_GET_SIZE(po->pool); 3299 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) { 3300 PyErr_SetString(PyExc_ValueError, "invalid arguments"); 3301 return NULL; 3302 } 3303 3304 for (i=0; i<n; i++) { 3305 PyObject* indexObject = PyTuple_GET_ITEM(indices, i); 3306 Py_ssize_t index = PyLong_AsSsize_t(indexObject); 3307 if (index < 0 && PyErr_Occurred()) 3308 return NULL; /* not an integer */ 3309 /* clamp the index */ 3310 if (index < 0) 3311 index = 0; 3312 else if (index > n-1) 3313 index = n-1; 3314 po->indices[i] = index; 3315 } 3316 3317 for (i=0; i<po->r; i++) { 3318 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i); 3319 Py_ssize_t index = PyLong_AsSsize_t(indexObject); 3320 if (index < 0 && PyErr_Occurred()) 3321 return NULL; /* not an integer */ 3322 if (index < 1) 3323 index = 1; 3324 else if (index > n-i) 3325 index = n-i; 3326 po->cycles[i] = index; 3327 } 3328 result = PyTuple_New(po->r); 3329 if (result == NULL) 3330 return NULL; 3331 for (i=0; i<po->r; i++) { 3332 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]); 3333 Py_INCREF(element); 3334 PyTuple_SET_ITEM(result, i, element); 3335 } 3336 Py_XSETREF(po->result, result); 3337 Py_RETURN_NONE; 3338 } 3339 3340 static PyMethodDef permuations_methods[] = { 3341 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS, 3342 reduce_doc}, 3343 {"__setstate__", (PyCFunction)permutations_setstate, METH_O, 3344 setstate_doc}, 3345 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS, 3346 sizeof_doc}, 3347 {NULL, NULL} /* sentinel */ 3348 }; 3349 3350 PyDoc_STRVAR(permutations_doc, 3351 "permutations(iterable[, r]) --> permutations object\n\ 3352 \n\ 3353 Return successive r-length permutations of elements in the iterable.\n\n\ 3354 permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)"); 3355 3356 static PyTypeObject permutations_type = { 3357 PyVarObject_HEAD_INIT(NULL, 0) 3358 "itertools.permutations", /* tp_name */ 3359 sizeof(permutationsobject), /* tp_basicsize */ 3360 0, /* tp_itemsize */ 3361 /* methods */ 3362 (destructor)permutations_dealloc, /* tp_dealloc */ 3363 0, /* tp_print */ 3364 0, /* tp_getattr */ 3365 0, /* tp_setattr */ 3366 0, /* tp_reserved */ 3367 0, /* tp_repr */ 3368 0, /* tp_as_number */ 3369 0, /* tp_as_sequence */ 3370 0, /* tp_as_mapping */ 3371 0, /* tp_hash */ 3372 0, /* tp_call */ 3373 0, /* tp_str */ 3374 PyObject_GenericGetAttr, /* tp_getattro */ 3375 0, /* tp_setattro */ 3376 0, /* tp_as_buffer */ 3377 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 3378 Py_TPFLAGS_BASETYPE, /* tp_flags */ 3379 permutations_doc, /* tp_doc */ 3380 (traverseproc)permutations_traverse,/* tp_traverse */ 3381 0, /* tp_clear */ 3382 0, /* tp_richcompare */ 3383 0, /* tp_weaklistoffset */ 3384 PyObject_SelfIter, /* tp_iter */ 3385 (iternextfunc)permutations_next, /* tp_iternext */ 3386 permuations_methods, /* tp_methods */ 3387 0, /* tp_members */ 3388 0, /* tp_getset */ 3389 0, /* tp_base */ 3390 0, /* tp_dict */ 3391 0, /* tp_descr_get */ 3392 0, /* tp_descr_set */ 3393 0, /* tp_dictoffset */ 3394 0, /* tp_init */ 3395 0, /* tp_alloc */ 3396 permutations_new, /* tp_new */ 3397 PyObject_GC_Del, /* tp_free */ 3398 }; 3399 3400 /* accumulate object ********************************************************/ 3401 3402 typedef struct { 3403 PyObject_HEAD 3404 PyObject *total; 3405 PyObject *it; 3406 PyObject *binop; 3407 } accumulateobject; 3408 3409 static PyTypeObject accumulate_type; 3410 3411 static PyObject * 3412 accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 3413 { 3414 static char *kwargs[] = {"iterable", "func", NULL}; 3415 PyObject *iterable; 3416 PyObject *it; 3417 PyObject *binop = Py_None; 3418 accumulateobject *lz; 3419 3420 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate", 3421 kwargs, &iterable, &binop)) 3422 return NULL; 3423 3424 /* Get iterator. */ 3425 it = PyObject_GetIter(iterable); 3426 if (it == NULL) 3427 return NULL; 3428 3429 /* create accumulateobject structure */ 3430 lz = (accumulateobject *)type->tp_alloc(type, 0); 3431 if (lz == NULL) { 3432 Py_DECREF(it); 3433 return NULL; 3434 } 3435 3436 if (binop != Py_None) { 3437 Py_XINCREF(binop); 3438 lz->binop = binop; 3439 } 3440 lz->total = NULL; 3441 lz->it = it; 3442 return (PyObject *)lz; 3443 } 3444 3445 static void 3446 accumulate_dealloc(accumulateobject *lz) 3447 { 3448 PyObject_GC_UnTrack(lz); 3449 Py_XDECREF(lz->binop); 3450 Py_XDECREF(lz->total); 3451 Py_XDECREF(lz->it); 3452 Py_TYPE(lz)->tp_free(lz); 3453 } 3454 3455 static int 3456 accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg) 3457 { 3458 Py_VISIT(lz->binop); 3459 Py_VISIT(lz->it); 3460 Py_VISIT(lz->total); 3461 return 0; 3462 } 3463 3464 static PyObject * 3465 accumulate_next(accumulateobject *lz) 3466 { 3467 PyObject *val, *newtotal; 3468 3469 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it); 3470 if (val == NULL) 3471 return NULL; 3472 3473 if (lz->total == NULL) { 3474 Py_INCREF(val); 3475 lz->total = val; 3476 return lz->total; 3477 } 3478 3479 if (lz->binop == NULL) 3480 newtotal = PyNumber_Add(lz->total, val); 3481 else 3482 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL); 3483 Py_DECREF(val); 3484 if (newtotal == NULL) 3485 return NULL; 3486 3487 Py_INCREF(newtotal); 3488 Py_SETREF(lz->total, newtotal); 3489 return newtotal; 3490 } 3491 3492 static PyObject * 3493 accumulate_reduce(accumulateobject *lz) 3494 { 3495 if (lz->total == Py_None) { 3496 PyObject *it; 3497 3498 if (PyType_Ready(&chain_type) < 0) 3499 return NULL; 3500 if (PyType_Ready(&islice_type) < 0) 3501 return NULL; 3502 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O", 3503 lz->total, lz->it); 3504 if (it == NULL) 3505 return NULL; 3506 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO", 3507 it, lz->binop ? lz->binop : Py_None); 3508 if (it == NULL) 3509 return NULL; 3510 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None); 3511 } 3512 return Py_BuildValue("O(OO)O", Py_TYPE(lz), 3513 lz->it, lz->binop?lz->binop:Py_None, 3514 lz->total?lz->total:Py_None); 3515 } 3516 3517 static PyObject * 3518 accumulate_setstate(accumulateobject *lz, PyObject *state) 3519 { 3520 Py_INCREF(state); 3521 Py_XSETREF(lz->total, state); 3522 Py_RETURN_NONE; 3523 } 3524 3525 static PyMethodDef accumulate_methods[] = { 3526 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS, 3527 reduce_doc}, 3528 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O, 3529 setstate_doc}, 3530 {NULL, NULL} /* sentinel */ 3531 }; 3532 3533 PyDoc_STRVAR(accumulate_doc, 3534 "accumulate(iterable[, func]) --> accumulate object\n\ 3535 \n\ 3536 Return series of accumulated sums (or other binary function results)."); 3537 3538 static PyTypeObject accumulate_type = { 3539 PyVarObject_HEAD_INIT(NULL, 0) 3540 "itertools.accumulate", /* tp_name */ 3541 sizeof(accumulateobject), /* tp_basicsize */ 3542 0, /* tp_itemsize */ 3543 /* methods */ 3544 (destructor)accumulate_dealloc, /* tp_dealloc */ 3545 0, /* tp_print */ 3546 0, /* tp_getattr */ 3547 0, /* tp_setattr */ 3548 0, /* tp_reserved */ 3549 0, /* tp_repr */ 3550 0, /* tp_as_number */ 3551 0, /* tp_as_sequence */ 3552 0, /* tp_as_mapping */ 3553 0, /* tp_hash */ 3554 0, /* tp_call */ 3555 0, /* tp_str */ 3556 PyObject_GenericGetAttr, /* tp_getattro */ 3557 0, /* tp_setattro */ 3558 0, /* tp_as_buffer */ 3559 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 3560 Py_TPFLAGS_BASETYPE, /* tp_flags */ 3561 accumulate_doc, /* tp_doc */ 3562 (traverseproc)accumulate_traverse, /* tp_traverse */ 3563 0, /* tp_clear */ 3564 0, /* tp_richcompare */ 3565 0, /* tp_weaklistoffset */ 3566 PyObject_SelfIter, /* tp_iter */ 3567 (iternextfunc)accumulate_next, /* tp_iternext */ 3568 accumulate_methods, /* tp_methods */ 3569 0, /* tp_members */ 3570 0, /* tp_getset */ 3571 0, /* tp_base */ 3572 0, /* tp_dict */ 3573 0, /* tp_descr_get */ 3574 0, /* tp_descr_set */ 3575 0, /* tp_dictoffset */ 3576 0, /* tp_init */ 3577 0, /* tp_alloc */ 3578 accumulate_new, /* tp_new */ 3579 PyObject_GC_Del, /* tp_free */ 3580 }; 3581 3582 3583 /* compress object ************************************************************/ 3584 3585 /* Equivalent to: 3586 3587 def compress(data, selectors): 3588 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F" 3589 return (d for d, s in zip(data, selectors) if s) 3590 */ 3591 3592 typedef struct { 3593 PyObject_HEAD 3594 PyObject *data; 3595 PyObject *selectors; 3596 } compressobject; 3597 3598 static PyTypeObject compress_type; 3599 3600 static PyObject * 3601 compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 3602 { 3603 PyObject *seq1, *seq2; 3604 PyObject *data=NULL, *selectors=NULL; 3605 compressobject *lz; 3606 static char *kwargs[] = {"data", "selectors", NULL}; 3607 3608 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2)) 3609 return NULL; 3610 3611 data = PyObject_GetIter(seq1); 3612 if (data == NULL) 3613 goto fail; 3614 selectors = PyObject_GetIter(seq2); 3615 if (selectors == NULL) 3616 goto fail; 3617 3618 /* create compressobject structure */ 3619 lz = (compressobject *)type->tp_alloc(type, 0); 3620 if (lz == NULL) 3621 goto fail; 3622 lz->data = data; 3623 lz->selectors = selectors; 3624 return (PyObject *)lz; 3625 3626 fail: 3627 Py_XDECREF(data); 3628 Py_XDECREF(selectors); 3629 return NULL; 3630 } 3631 3632 static void 3633 compress_dealloc(compressobject *lz) 3634 { 3635 PyObject_GC_UnTrack(lz); 3636 Py_XDECREF(lz->data); 3637 Py_XDECREF(lz->selectors); 3638 Py_TYPE(lz)->tp_free(lz); 3639 } 3640 3641 static int 3642 compress_traverse(compressobject *lz, visitproc visit, void *arg) 3643 { 3644 Py_VISIT(lz->data); 3645 Py_VISIT(lz->selectors); 3646 return 0; 3647 } 3648 3649 static PyObject * 3650 compress_next(compressobject *lz) 3651 { 3652 PyObject *data = lz->data, *selectors = lz->selectors; 3653 PyObject *datum, *selector; 3654 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext; 3655 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext; 3656 int ok; 3657 3658 while (1) { 3659 /* Steps: get datum, get selector, evaluate selector. 3660 Order is important (to match the pure python version 3661 in terms of which input gets a chance to raise an 3662 exception first). 3663 */ 3664 3665 datum = datanext(data); 3666 if (datum == NULL) 3667 return NULL; 3668 3669 selector = selectornext(selectors); 3670 if (selector == NULL) { 3671 Py_DECREF(datum); 3672 return NULL; 3673 } 3674 3675 ok = PyObject_IsTrue(selector); 3676 Py_DECREF(selector); 3677 if (ok > 0) 3678 return datum; 3679 Py_DECREF(datum); 3680 if (ok < 0) 3681 return NULL; 3682 } 3683 } 3684 3685 static PyObject * 3686 compress_reduce(compressobject *lz) 3687 { 3688 return Py_BuildValue("O(OO)", Py_TYPE(lz), 3689 lz->data, lz->selectors); 3690 } 3691 3692 static PyMethodDef compress_methods[] = { 3693 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS, 3694 reduce_doc}, 3695 {NULL, NULL} /* sentinel */ 3696 }; 3697 3698 PyDoc_STRVAR(compress_doc, 3699 "compress(data, selectors) --> iterator over selected data\n\ 3700 \n\ 3701 Return data elements corresponding to true selector elements.\n\ 3702 Forms a shorter iterator from selected data elements using the\n\ 3703 selectors to choose the data elements."); 3704 3705 static PyTypeObject compress_type = { 3706 PyVarObject_HEAD_INIT(NULL, 0) 3707 "itertools.compress", /* tp_name */ 3708 sizeof(compressobject), /* tp_basicsize */ 3709 0, /* tp_itemsize */ 3710 /* methods */ 3711 (destructor)compress_dealloc, /* tp_dealloc */ 3712 0, /* tp_print */ 3713 0, /* tp_getattr */ 3714 0, /* tp_setattr */ 3715 0, /* tp_reserved */ 3716 0, /* tp_repr */ 3717 0, /* tp_as_number */ 3718 0, /* tp_as_sequence */ 3719 0, /* tp_as_mapping */ 3720 0, /* tp_hash */ 3721 0, /* tp_call */ 3722 0, /* tp_str */ 3723 PyObject_GenericGetAttr, /* tp_getattro */ 3724 0, /* tp_setattro */ 3725 0, /* tp_as_buffer */ 3726 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 3727 Py_TPFLAGS_BASETYPE, /* tp_flags */ 3728 compress_doc, /* tp_doc */ 3729 (traverseproc)compress_traverse, /* tp_traverse */ 3730 0, /* tp_clear */ 3731 0, /* tp_richcompare */ 3732 0, /* tp_weaklistoffset */ 3733 PyObject_SelfIter, /* tp_iter */ 3734 (iternextfunc)compress_next, /* tp_iternext */ 3735 compress_methods, /* tp_methods */ 3736 0, /* tp_members */ 3737 0, /* tp_getset */ 3738 0, /* tp_base */ 3739 0, /* tp_dict */ 3740 0, /* tp_descr_get */ 3741 0, /* tp_descr_set */ 3742 0, /* tp_dictoffset */ 3743 0, /* tp_init */ 3744 0, /* tp_alloc */ 3745 compress_new, /* tp_new */ 3746 PyObject_GC_Del, /* tp_free */ 3747 }; 3748 3749 3750 /* filterfalse object ************************************************************/ 3751 3752 typedef struct { 3753 PyObject_HEAD 3754 PyObject *func; 3755 PyObject *it; 3756 } filterfalseobject; 3757 3758 static PyTypeObject filterfalse_type; 3759 3760 static PyObject * 3761 filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 3762 { 3763 PyObject *func, *seq; 3764 PyObject *it; 3765 filterfalseobject *lz; 3766 3767 if (type == &filterfalse_type && 3768 !_PyArg_NoKeywords("filterfalse()", kwds)) 3769 return NULL; 3770 3771 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq)) 3772 return NULL; 3773 3774 /* Get iterator. */ 3775 it = PyObject_GetIter(seq); 3776 if (it == NULL) 3777 return NULL; 3778 3779 /* create filterfalseobject structure */ 3780 lz = (filterfalseobject *)type->tp_alloc(type, 0); 3781 if (lz == NULL) { 3782 Py_DECREF(it); 3783 return NULL; 3784 } 3785 Py_INCREF(func); 3786 lz->func = func; 3787 lz->it = it; 3788 3789 return (PyObject *)lz; 3790 } 3791 3792 static void 3793 filterfalse_dealloc(filterfalseobject *lz) 3794 { 3795 PyObject_GC_UnTrack(lz); 3796 Py_XDECREF(lz->func); 3797 Py_XDECREF(lz->it); 3798 Py_TYPE(lz)->tp_free(lz); 3799 } 3800 3801 static int 3802 filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg) 3803 { 3804 Py_VISIT(lz->it); 3805 Py_VISIT(lz->func); 3806 return 0; 3807 } 3808 3809 static PyObject * 3810 filterfalse_next(filterfalseobject *lz) 3811 { 3812 PyObject *item; 3813 PyObject *it = lz->it; 3814 long ok; 3815 PyObject *(*iternext)(PyObject *); 3816 3817 iternext = *Py_TYPE(it)->tp_iternext; 3818 for (;;) { 3819 item = iternext(it); 3820 if (item == NULL) 3821 return NULL; 3822 3823 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) { 3824 ok = PyObject_IsTrue(item); 3825 } else { 3826 PyObject *good; 3827 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL); 3828 if (good == NULL) { 3829 Py_DECREF(item); 3830 return NULL; 3831 } 3832 ok = PyObject_IsTrue(good); 3833 Py_DECREF(good); 3834 } 3835 if (ok == 0) 3836 return item; 3837 Py_DECREF(item); 3838 if (ok < 0) 3839 return NULL; 3840 } 3841 } 3842 3843 static PyObject * 3844 filterfalse_reduce(filterfalseobject *lz) 3845 { 3846 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it); 3847 } 3848 3849 static PyMethodDef filterfalse_methods[] = { 3850 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS, 3851 reduce_doc}, 3852 {NULL, NULL} /* sentinel */ 3853 }; 3854 3855 PyDoc_STRVAR(filterfalse_doc, 3856 "filterfalse(function or None, sequence) --> filterfalse object\n\ 3857 \n\ 3858 Return those items of sequence for which function(item) is false.\n\ 3859 If function is None, return the items that are false."); 3860 3861 static PyTypeObject filterfalse_type = { 3862 PyVarObject_HEAD_INIT(NULL, 0) 3863 "itertools.filterfalse", /* tp_name */ 3864 sizeof(filterfalseobject), /* tp_basicsize */ 3865 0, /* tp_itemsize */ 3866 /* methods */ 3867 (destructor)filterfalse_dealloc, /* tp_dealloc */ 3868 0, /* tp_print */ 3869 0, /* tp_getattr */ 3870 0, /* tp_setattr */ 3871 0, /* tp_reserved */ 3872 0, /* tp_repr */ 3873 0, /* tp_as_number */ 3874 0, /* tp_as_sequence */ 3875 0, /* tp_as_mapping */ 3876 0, /* tp_hash */ 3877 0, /* tp_call */ 3878 0, /* tp_str */ 3879 PyObject_GenericGetAttr, /* tp_getattro */ 3880 0, /* tp_setattro */ 3881 0, /* tp_as_buffer */ 3882 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 3883 Py_TPFLAGS_BASETYPE, /* tp_flags */ 3884 filterfalse_doc, /* tp_doc */ 3885 (traverseproc)filterfalse_traverse, /* tp_traverse */ 3886 0, /* tp_clear */ 3887 0, /* tp_richcompare */ 3888 0, /* tp_weaklistoffset */ 3889 PyObject_SelfIter, /* tp_iter */ 3890 (iternextfunc)filterfalse_next, /* tp_iternext */ 3891 filterfalse_methods, /* tp_methods */ 3892 0, /* tp_members */ 3893 0, /* tp_getset */ 3894 0, /* tp_base */ 3895 0, /* tp_dict */ 3896 0, /* tp_descr_get */ 3897 0, /* tp_descr_set */ 3898 0, /* tp_dictoffset */ 3899 0, /* tp_init */ 3900 0, /* tp_alloc */ 3901 filterfalse_new, /* tp_new */ 3902 PyObject_GC_Del, /* tp_free */ 3903 }; 3904 3905 3906 /* count object ************************************************************/ 3907 3908 typedef struct { 3909 PyObject_HEAD 3910 Py_ssize_t cnt; 3911 PyObject *long_cnt; 3912 PyObject *long_step; 3913 } countobject; 3914 3915 /* Counting logic and invariants: 3916 3917 fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified. 3918 3919 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1)); 3920 Advances with: cnt += 1 3921 When count hits Y_SSIZE_T_MAX, switch to slow_mode. 3922 3923 slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float. 3924 3925 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL); 3926 All counting is done with python objects (no overflows or underflows). 3927 Advances with: long_cnt += long_step 3928 Step may be zero -- effectively a slow version of repeat(cnt). 3929 Either long_cnt or long_step may be a float, Fraction, or Decimal. 3930 */ 3931 3932 static PyTypeObject count_type; 3933 3934 static PyObject * 3935 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 3936 { 3937 countobject *lz; 3938 int fast_mode; 3939 Py_ssize_t cnt = 0; 3940 PyObject *long_cnt = NULL; 3941 PyObject *long_step = NULL; 3942 long step; 3943 static char *kwlist[] = {"start", "step", 0}; 3944 3945 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count", 3946 kwlist, &long_cnt, &long_step)) 3947 return NULL; 3948 3949 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) || 3950 (long_step != NULL && !PyNumber_Check(long_step))) { 3951 PyErr_SetString(PyExc_TypeError, "a number is required"); 3952 return NULL; 3953 } 3954 3955 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) && 3956 (long_step == NULL || PyLong_Check(long_step)); 3957 3958 /* If not specified, start defaults to 0 */ 3959 if (long_cnt != NULL) { 3960 if (fast_mode) { 3961 assert(PyLong_Check(long_cnt)); 3962 cnt = PyLong_AsSsize_t(long_cnt); 3963 if (cnt == -1 && PyErr_Occurred()) { 3964 PyErr_Clear(); 3965 fast_mode = 0; 3966 } 3967 } 3968 Py_INCREF(long_cnt); 3969 } else { 3970 cnt = 0; 3971 long_cnt = PyLong_FromLong(0); 3972 if (long_cnt == NULL) { 3973 return NULL; 3974 } 3975 } 3976 3977 /* If not specified, step defaults to 1 */ 3978 if (long_step == NULL) { 3979 long_step = PyLong_FromLong(1); 3980 if (long_step == NULL) { 3981 Py_DECREF(long_cnt); 3982 return NULL; 3983 } 3984 } else 3985 Py_INCREF(long_step); 3986 3987 assert(long_cnt != NULL && long_step != NULL); 3988 3989 /* Fast mode only works when the step is 1 */ 3990 if (fast_mode) { 3991 assert(PyLong_Check(long_step)); 3992 step = PyLong_AsLong(long_step); 3993 if (step != 1) { 3994 fast_mode = 0; 3995 if (step == -1 && PyErr_Occurred()) 3996 PyErr_Clear(); 3997 } 3998 } 3999 4000 if (fast_mode) 4001 Py_CLEAR(long_cnt); 4002 else 4003 cnt = PY_SSIZE_T_MAX; 4004 4005 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) || 4006 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode)); 4007 assert(!fast_mode || 4008 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1)); 4009 4010 /* create countobject structure */ 4011 lz = (countobject *)type->tp_alloc(type, 0); 4012 if (lz == NULL) { 4013 Py_XDECREF(long_cnt); 4014 return NULL; 4015 } 4016 lz->cnt = cnt; 4017 lz->long_cnt = long_cnt; 4018 lz->long_step = long_step; 4019 4020 return (PyObject *)lz; 4021 } 4022 4023 static void 4024 count_dealloc(countobject *lz) 4025 { 4026 PyObject_GC_UnTrack(lz); 4027 Py_XDECREF(lz->long_cnt); 4028 Py_XDECREF(lz->long_step); 4029 Py_TYPE(lz)->tp_free(lz); 4030 } 4031 4032 static int 4033 count_traverse(countobject *lz, visitproc visit, void *arg) 4034 { 4035 Py_VISIT(lz->long_cnt); 4036 Py_VISIT(lz->long_step); 4037 return 0; 4038 } 4039 4040 static PyObject * 4041 count_nextlong(countobject *lz) 4042 { 4043 PyObject *long_cnt; 4044 PyObject *stepped_up; 4045 4046 long_cnt = lz->long_cnt; 4047 if (long_cnt == NULL) { 4048 /* Switch to slow_mode */ 4049 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX); 4050 if (long_cnt == NULL) 4051 return NULL; 4052 } 4053 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL); 4054 4055 stepped_up = PyNumber_Add(long_cnt, lz->long_step); 4056 if (stepped_up == NULL) 4057 return NULL; 4058 lz->long_cnt = stepped_up; 4059 return long_cnt; 4060 } 4061 4062 static PyObject * 4063 count_next(countobject *lz) 4064 { 4065 if (lz->cnt == PY_SSIZE_T_MAX) 4066 return count_nextlong(lz); 4067 return PyLong_FromSsize_t(lz->cnt++); 4068 } 4069 4070 static PyObject * 4071 count_repr(countobject *lz) 4072 { 4073 if (lz->cnt != PY_SSIZE_T_MAX) 4074 return PyUnicode_FromFormat("count(%zd)", lz->cnt); 4075 4076 if (PyLong_Check(lz->long_step)) { 4077 long step = PyLong_AsLong(lz->long_step); 4078 if (step == -1 && PyErr_Occurred()) { 4079 PyErr_Clear(); 4080 } 4081 if (step == 1) { 4082 /* Don't display step when it is an integer equal to 1 */ 4083 return PyUnicode_FromFormat("count(%R)", lz->long_cnt); 4084 } 4085 } 4086 return PyUnicode_FromFormat("count(%R, %R)", 4087 lz->long_cnt, lz->long_step); 4088 } 4089 4090 static PyObject * 4091 count_reduce(countobject *lz) 4092 { 4093 if (lz->cnt == PY_SSIZE_T_MAX) 4094 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step); 4095 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt); 4096 } 4097 4098 static PyMethodDef count_methods[] = { 4099 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS, 4100 reduce_doc}, 4101 {NULL, NULL} /* sentinel */ 4102 }; 4103 4104 PyDoc_STRVAR(count_doc, 4105 "count(start=0, step=1) --> count object\n\ 4106 \n\ 4107 Return a count object whose .__next__() method returns consecutive values.\n\ 4108 Equivalent to:\n\n\ 4109 def count(firstval=0, step=1):\n\ 4110 x = firstval\n\ 4111 while 1:\n\ 4112 yield x\n\ 4113 x += step\n"); 4114 4115 static PyTypeObject count_type = { 4116 PyVarObject_HEAD_INIT(NULL, 0) 4117 "itertools.count", /* tp_name */ 4118 sizeof(countobject), /* tp_basicsize */ 4119 0, /* tp_itemsize */ 4120 /* methods */ 4121 (destructor)count_dealloc, /* tp_dealloc */ 4122 0, /* tp_print */ 4123 0, /* tp_getattr */ 4124 0, /* tp_setattr */ 4125 0, /* tp_reserved */ 4126 (reprfunc)count_repr, /* tp_repr */ 4127 0, /* tp_as_number */ 4128 0, /* tp_as_sequence */ 4129 0, /* tp_as_mapping */ 4130 0, /* tp_hash */ 4131 0, /* tp_call */ 4132 0, /* tp_str */ 4133 PyObject_GenericGetAttr, /* tp_getattro */ 4134 0, /* tp_setattro */ 4135 0, /* tp_as_buffer */ 4136 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 4137 Py_TPFLAGS_BASETYPE, /* tp_flags */ 4138 count_doc, /* tp_doc */ 4139 (traverseproc)count_traverse, /* tp_traverse */ 4140 0, /* tp_clear */ 4141 0, /* tp_richcompare */ 4142 0, /* tp_weaklistoffset */ 4143 PyObject_SelfIter, /* tp_iter */ 4144 (iternextfunc)count_next, /* tp_iternext */ 4145 count_methods, /* tp_methods */ 4146 0, /* tp_members */ 4147 0, /* tp_getset */ 4148 0, /* tp_base */ 4149 0, /* tp_dict */ 4150 0, /* tp_descr_get */ 4151 0, /* tp_descr_set */ 4152 0, /* tp_dictoffset */ 4153 0, /* tp_init */ 4154 0, /* tp_alloc */ 4155 count_new, /* tp_new */ 4156 PyObject_GC_Del, /* tp_free */ 4157 }; 4158 4159 4160 /* repeat object ************************************************************/ 4161 4162 typedef struct { 4163 PyObject_HEAD 4164 PyObject *element; 4165 Py_ssize_t cnt; 4166 } repeatobject; 4167 4168 static PyTypeObject repeat_type; 4169 4170 static PyObject * 4171 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 4172 { 4173 repeatobject *ro; 4174 PyObject *element; 4175 Py_ssize_t cnt = -1, n_kwds = 0; 4176 static char *kwargs[] = {"object", "times", NULL}; 4177 4178 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs, 4179 &element, &cnt)) 4180 return NULL; 4181 4182 if (kwds != NULL) 4183 n_kwds = PyDict_Size(kwds); 4184 /* Does user supply times argument? */ 4185 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0) 4186 cnt = 0; 4187 4188 ro = (repeatobject *)type->tp_alloc(type, 0); 4189 if (ro == NULL) 4190 return NULL; 4191 Py_INCREF(element); 4192 ro->element = element; 4193 ro->cnt = cnt; 4194 return (PyObject *)ro; 4195 } 4196 4197 static void 4198 repeat_dealloc(repeatobject *ro) 4199 { 4200 PyObject_GC_UnTrack(ro); 4201 Py_XDECREF(ro->element); 4202 Py_TYPE(ro)->tp_free(ro); 4203 } 4204 4205 static int 4206 repeat_traverse(repeatobject *ro, visitproc visit, void *arg) 4207 { 4208 Py_VISIT(ro->element); 4209 return 0; 4210 } 4211 4212 static PyObject * 4213 repeat_next(repeatobject *ro) 4214 { 4215 if (ro->cnt == 0) 4216 return NULL; 4217 if (ro->cnt > 0) 4218 ro->cnt--; 4219 Py_INCREF(ro->element); 4220 return ro->element; 4221 } 4222 4223 static PyObject * 4224 repeat_repr(repeatobject *ro) 4225 { 4226 if (ro->cnt == -1) 4227 return PyUnicode_FromFormat("repeat(%R)", ro->element); 4228 else 4229 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt); 4230 } 4231 4232 static PyObject * 4233 repeat_len(repeatobject *ro) 4234 { 4235 if (ro->cnt == -1) { 4236 PyErr_SetString(PyExc_TypeError, "len() of unsized object"); 4237 return NULL; 4238 } 4239 return PyLong_FromSize_t(ro->cnt); 4240 } 4241 4242 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it))."); 4243 4244 static PyObject * 4245 repeat_reduce(repeatobject *ro) 4246 { 4247 /* unpickle this so that a new repeat iterator is constructed with an 4248 * object, then call __setstate__ on it to set cnt 4249 */ 4250 if (ro->cnt >= 0) 4251 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt); 4252 else 4253 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element); 4254 } 4255 4256 static PyMethodDef repeat_methods[] = { 4257 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc}, 4258 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc}, 4259 {NULL, NULL} /* sentinel */ 4260 }; 4261 4262 PyDoc_STRVAR(repeat_doc, 4263 "repeat(object [,times]) -> create an iterator which returns the object\n\ 4264 for the specified number of times. If not specified, returns the object\n\ 4265 endlessly."); 4266 4267 static PyTypeObject repeat_type = { 4268 PyVarObject_HEAD_INIT(NULL, 0) 4269 "itertools.repeat", /* tp_name */ 4270 sizeof(repeatobject), /* tp_basicsize */ 4271 0, /* tp_itemsize */ 4272 /* methods */ 4273 (destructor)repeat_dealloc, /* tp_dealloc */ 4274 0, /* tp_print */ 4275 0, /* tp_getattr */ 4276 0, /* tp_setattr */ 4277 0, /* tp_reserved */ 4278 (reprfunc)repeat_repr, /* tp_repr */ 4279 0, /* tp_as_number */ 4280 0, /* tp_as_sequence */ 4281 0, /* tp_as_mapping */ 4282 0, /* tp_hash */ 4283 0, /* tp_call */ 4284 0, /* tp_str */ 4285 PyObject_GenericGetAttr, /* tp_getattro */ 4286 0, /* tp_setattro */ 4287 0, /* tp_as_buffer */ 4288 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 4289 Py_TPFLAGS_BASETYPE, /* tp_flags */ 4290 repeat_doc, /* tp_doc */ 4291 (traverseproc)repeat_traverse, /* tp_traverse */ 4292 0, /* tp_clear */ 4293 0, /* tp_richcompare */ 4294 0, /* tp_weaklistoffset */ 4295 PyObject_SelfIter, /* tp_iter */ 4296 (iternextfunc)repeat_next, /* tp_iternext */ 4297 repeat_methods, /* tp_methods */ 4298 0, /* tp_members */ 4299 0, /* tp_getset */ 4300 0, /* tp_base */ 4301 0, /* tp_dict */ 4302 0, /* tp_descr_get */ 4303 0, /* tp_descr_set */ 4304 0, /* tp_dictoffset */ 4305 0, /* tp_init */ 4306 0, /* tp_alloc */ 4307 repeat_new, /* tp_new */ 4308 PyObject_GC_Del, /* tp_free */ 4309 }; 4310 4311 /* ziplongest object *********************************************************/ 4312 4313 typedef struct { 4314 PyObject_HEAD 4315 Py_ssize_t tuplesize; 4316 Py_ssize_t numactive; 4317 PyObject *ittuple; /* tuple of iterators */ 4318 PyObject *result; 4319 PyObject *fillvalue; 4320 } ziplongestobject; 4321 4322 static PyTypeObject ziplongest_type; 4323 4324 static PyObject * 4325 zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds) 4326 { 4327 ziplongestobject *lz; 4328 Py_ssize_t i; 4329 PyObject *ittuple; /* tuple of iterators */ 4330 PyObject *result; 4331 PyObject *fillvalue = Py_None; 4332 Py_ssize_t tuplesize = PySequence_Length(args); 4333 4334 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) { 4335 fillvalue = PyDict_GetItemString(kwds, "fillvalue"); 4336 if (fillvalue == NULL || PyDict_Size(kwds) > 1) { 4337 PyErr_SetString(PyExc_TypeError, 4338 "zip_longest() got an unexpected keyword argument"); 4339 return NULL; 4340 } 4341 } 4342 4343 /* args must be a tuple */ 4344 assert(PyTuple_Check(args)); 4345 4346 /* obtain iterators */ 4347 ittuple = PyTuple_New(tuplesize); 4348 if (ittuple == NULL) 4349 return NULL; 4350 for (i=0; i < tuplesize; i++) { 4351 PyObject *item = PyTuple_GET_ITEM(args, i); 4352 PyObject *it = PyObject_GetIter(item); 4353 if (it == NULL) { 4354 if (PyErr_ExceptionMatches(PyExc_TypeError)) 4355 PyErr_Format(PyExc_TypeError, 4356 "zip_longest argument #%zd must support iteration", 4357 i+1); 4358 Py_DECREF(ittuple); 4359 return NULL; 4360 } 4361 PyTuple_SET_ITEM(ittuple, i, it); 4362 } 4363 4364 /* create a result holder */ 4365 result = PyTuple_New(tuplesize); 4366 if (result == NULL) { 4367 Py_DECREF(ittuple); 4368 return NULL; 4369 } 4370 for (i=0 ; i < tuplesize ; i++) { 4371 Py_INCREF(Py_None); 4372 PyTuple_SET_ITEM(result, i, Py_None); 4373 } 4374 4375 /* create ziplongestobject structure */ 4376 lz = (ziplongestobject *)type->tp_alloc(type, 0); 4377 if (lz == NULL) { 4378 Py_DECREF(ittuple); 4379 Py_DECREF(result); 4380 return NULL; 4381 } 4382 lz->ittuple = ittuple; 4383 lz->tuplesize = tuplesize; 4384 lz->numactive = tuplesize; 4385 lz->result = result; 4386 Py_INCREF(fillvalue); 4387 lz->fillvalue = fillvalue; 4388 return (PyObject *)lz; 4389 } 4390 4391 static void 4392 zip_longest_dealloc(ziplongestobject *lz) 4393 { 4394 PyObject_GC_UnTrack(lz); 4395 Py_XDECREF(lz->ittuple); 4396 Py_XDECREF(lz->result); 4397 Py_XDECREF(lz->fillvalue); 4398 Py_TYPE(lz)->tp_free(lz); 4399 } 4400 4401 static int 4402 zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg) 4403 { 4404 Py_VISIT(lz->ittuple); 4405 Py_VISIT(lz->result); 4406 Py_VISIT(lz->fillvalue); 4407 return 0; 4408 } 4409 4410 static PyObject * 4411 zip_longest_next(ziplongestobject *lz) 4412 { 4413 Py_ssize_t i; 4414 Py_ssize_t tuplesize = lz->tuplesize; 4415 PyObject *result = lz->result; 4416 PyObject *it; 4417 PyObject *item; 4418 PyObject *olditem; 4419 4420 if (tuplesize == 0) 4421 return NULL; 4422 if (lz->numactive == 0) 4423 return NULL; 4424 if (Py_REFCNT(result) == 1) { 4425 Py_INCREF(result); 4426 for (i=0 ; i < tuplesize ; i++) { 4427 it = PyTuple_GET_ITEM(lz->ittuple, i); 4428 if (it == NULL) { 4429 Py_INCREF(lz->fillvalue); 4430 item = lz->fillvalue; 4431 } else { 4432 item = PyIter_Next(it); 4433 if (item == NULL) { 4434 lz->numactive -= 1; 4435 if (lz->numactive == 0 || PyErr_Occurred()) { 4436 lz->numactive = 0; 4437 Py_DECREF(result); 4438 return NULL; 4439 } else { 4440 Py_INCREF(lz->fillvalue); 4441 item = lz->fillvalue; 4442 PyTuple_SET_ITEM(lz->ittuple, i, NULL); 4443 Py_DECREF(it); 4444 } 4445 } 4446 } 4447 olditem = PyTuple_GET_ITEM(result, i); 4448 PyTuple_SET_ITEM(result, i, item); 4449 Py_DECREF(olditem); 4450 } 4451 } else { 4452 result = PyTuple_New(tuplesize); 4453 if (result == NULL) 4454 return NULL; 4455 for (i=0 ; i < tuplesize ; i++) { 4456 it = PyTuple_GET_ITEM(lz->ittuple, i); 4457 if (it == NULL) { 4458 Py_INCREF(lz->fillvalue); 4459 item = lz->fillvalue; 4460 } else { 4461 item = PyIter_Next(it); 4462 if (item == NULL) { 4463 lz->numactive -= 1; 4464 if (lz->numactive == 0 || PyErr_Occurred()) { 4465 lz->numactive = 0; 4466 Py_DECREF(result); 4467 return NULL; 4468 } else { 4469 Py_INCREF(lz->fillvalue); 4470 item = lz->fillvalue; 4471 PyTuple_SET_ITEM(lz->ittuple, i, NULL); 4472 Py_DECREF(it); 4473 } 4474 } 4475 } 4476 PyTuple_SET_ITEM(result, i, item); 4477 } 4478 } 4479 return result; 4480 } 4481 4482 static PyObject * 4483 zip_longest_reduce(ziplongestobject *lz) 4484 { 4485 4486 /* Create a new tuple with empty sequences where appropriate to pickle. 4487 * Then use setstate to set the fillvalue 4488 */ 4489 int i; 4490 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple)); 4491 4492 if (args == NULL) 4493 return NULL; 4494 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) { 4495 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i); 4496 if (elem == NULL) { 4497 elem = PyTuple_New(0); 4498 if (elem == NULL) { 4499 Py_DECREF(args); 4500 return NULL; 4501 } 4502 } else 4503 Py_INCREF(elem); 4504 PyTuple_SET_ITEM(args, i, elem); 4505 } 4506 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue); 4507 } 4508 4509 static PyObject * 4510 zip_longest_setstate(ziplongestobject *lz, PyObject *state) 4511 { 4512 Py_INCREF(state); 4513 Py_XSETREF(lz->fillvalue, state); 4514 Py_RETURN_NONE; 4515 } 4516 4517 static PyMethodDef zip_longest_methods[] = { 4518 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS, 4519 reduce_doc}, 4520 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O, 4521 setstate_doc}, 4522 {NULL, NULL} /* sentinel */ 4523 }; 4524 4525 PyDoc_STRVAR(zip_longest_doc, 4526 "zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\ 4527 \n\ 4528 Return a zip_longest object whose .__next__() method returns a tuple where\n\ 4529 the i-th element comes from the i-th iterable argument. The .__next__()\n\ 4530 method continues until the longest iterable in the argument sequence\n\ 4531 is exhausted and then it raises StopIteration. When the shorter iterables\n\ 4532 are exhausted, the fillvalue is substituted in their place. The fillvalue\n\ 4533 defaults to None or can be specified by a keyword argument.\n\ 4534 "); 4535 4536 static PyTypeObject ziplongest_type = { 4537 PyVarObject_HEAD_INIT(NULL, 0) 4538 "itertools.zip_longest", /* tp_name */ 4539 sizeof(ziplongestobject), /* tp_basicsize */ 4540 0, /* tp_itemsize */ 4541 /* methods */ 4542 (destructor)zip_longest_dealloc, /* tp_dealloc */ 4543 0, /* tp_print */ 4544 0, /* tp_getattr */ 4545 0, /* tp_setattr */ 4546 0, /* tp_reserved */ 4547 0, /* tp_repr */ 4548 0, /* tp_as_number */ 4549 0, /* tp_as_sequence */ 4550 0, /* tp_as_mapping */ 4551 0, /* tp_hash */ 4552 0, /* tp_call */ 4553 0, /* tp_str */ 4554 PyObject_GenericGetAttr, /* tp_getattro */ 4555 0, /* tp_setattro */ 4556 0, /* tp_as_buffer */ 4557 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | 4558 Py_TPFLAGS_BASETYPE, /* tp_flags */ 4559 zip_longest_doc, /* tp_doc */ 4560 (traverseproc)zip_longest_traverse, /* tp_traverse */ 4561 0, /* tp_clear */ 4562 0, /* tp_richcompare */ 4563 0, /* tp_weaklistoffset */ 4564 PyObject_SelfIter, /* tp_iter */ 4565 (iternextfunc)zip_longest_next, /* tp_iternext */ 4566 zip_longest_methods, /* tp_methods */ 4567 0, /* tp_members */ 4568 0, /* tp_getset */ 4569 0, /* tp_base */ 4570 0, /* tp_dict */ 4571 0, /* tp_descr_get */ 4572 0, /* tp_descr_set */ 4573 0, /* tp_dictoffset */ 4574 0, /* tp_init */ 4575 0, /* tp_alloc */ 4576 zip_longest_new, /* tp_new */ 4577 PyObject_GC_Del, /* tp_free */ 4578 }; 4579 4580 /* module level code ********************************************************/ 4581 4582 PyDoc_STRVAR(module_doc, 4583 "Functional tools for creating and using iterators.\n\ 4584 \n\ 4585 Infinite iterators:\n\ 4586 count(start=0, step=1) --> start, start+step, start+2*step, ...\n\ 4587 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\ 4588 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\ 4589 \n\ 4590 Iterators terminating on the shortest input sequence:\n\ 4591 accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\ 4592 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\ 4593 chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\ 4594 compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\ 4595 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\ 4596 groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\ 4597 filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\ 4598 islice(seq, [start,] stop [, step]) --> elements from\n\ 4599 seq[start:stop:step]\n\ 4600 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\ 4601 tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\ 4602 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\ 4603 zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\ 4604 \n\ 4605 Combinatoric generators:\n\ 4606 product(p, q, ... [repeat=1]) --> cartesian product\n\ 4607 permutations(p[, r])\n\ 4608 combinations(p, r)\n\ 4609 combinations_with_replacement(p, r)\n\ 4610 "); 4611 4612 4613 static PyMethodDef module_methods[] = { 4614 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc}, 4615 {NULL, NULL} /* sentinel */ 4616 }; 4617 4618 4619 static struct PyModuleDef itertoolsmodule = { 4620 PyModuleDef_HEAD_INIT, 4621 "itertools", 4622 module_doc, 4623 -1, 4624 module_methods, 4625 NULL, 4626 NULL, 4627 NULL, 4628 NULL 4629 }; 4630 4631 PyMODINIT_FUNC 4632 PyInit_itertools(void) 4633 { 4634 int i; 4635 PyObject *m; 4636 char *name; 4637 PyTypeObject *typelist[] = { 4638 &accumulate_type, 4639 &combinations_type, 4640 &cwr_type, 4641 &cycle_type, 4642 &dropwhile_type, 4643 &takewhile_type, 4644 &islice_type, 4645 &starmap_type, 4646 &chain_type, 4647 &compress_type, 4648 &filterfalse_type, 4649 &count_type, 4650 &ziplongest_type, 4651 &permutations_type, 4652 &product_type, 4653 &repeat_type, 4654 &groupby_type, 4655 &_grouper_type, 4656 &tee_type, 4657 &teedataobject_type, 4658 NULL 4659 }; 4660 4661 Py_TYPE(&teedataobject_type) = &PyType_Type; 4662 m = PyModule_Create(&itertoolsmodule); 4663 if (m == NULL) 4664 return NULL; 4665 4666 for (i=0 ; typelist[i] != NULL ; i++) { 4667 if (PyType_Ready(typelist[i]) < 0) 4668 return NULL; 4669 name = strchr(typelist[i]->tp_name, '.'); 4670 assert (name != NULL); 4671 Py_INCREF(typelist[i]); 4672 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]); 4673 } 4674 4675 return m; 4676 } 4677
