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