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