Modules/itertoolsmodule.c

   1
   2 #include "Python.h"
   3
   4 /* Itertools module written and maintained
   5    by Raymond D. Hettinger <python@rcn.com>
   6    Copyright (c) 2003 Python Software Foundation.
   7    All rights reserved.
   8 */
   9
  10 /* cycle object **********************************************************/
  11
  12 typedef struct {
  13         PyObject_HEAD
  14         PyObject *it;
  15         PyObject *saved;
  16         int firstpass;
  17 } cycleobject;
  18
  19 static PyTypeObject cycle_type;
  20
  21 static PyObject *
  22 cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  23 {
  24         PyObject *it;
  25         PyObject *iterable;
  26         PyObject *saved;
  27         cycleobject *lz;
  28
  29         if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
  30                 return NULL;
  31
  32         /* Get iterator. */
  33         it = PyObject_GetIter(iterable);
  34         if (it == NULL)
  35                 return NULL;
  36
  37         saved = PyList_New(0);
  38         if (saved == NULL) {
  39                 Py_DECREF(it);
  40                 return NULL;
  41         }
  42
  43         /* create cycleobject structure */
  44         lz = (cycleobject *)type->tp_alloc(type, 0);
  45         if (lz == NULL) {
  46                 Py_DECREF(it);
  47                 Py_DECREF(saved);
  48                 return NULL;
  49         }
  50         lz->it = it;
  51         lz->saved = saved;
  52         lz->firstpass = 0;
  53
  54         return (PyObject *)lz;
  55 }
  56
  57 static void
  58 cycle_dealloc(cycleobject *lz)
  59 {
  60         PyObject_GC_UnTrack(lz);
  61         Py_XDECREF(lz->saved);
  62         Py_XDECREF(lz->it);
  63         lz->ob_type->tp_free(lz);
  64 }
  65
  66 static int
  67 cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
  68 {
  69         int err;
  70
  71         if (lz->it) {
  72                 err = visit(lz->it, arg);
  73                 if (err)
  74                         return err;
  75         }
  76         if (lz->saved) {
  77                 err = visit(lz->saved, arg);
  78                 if (err)
  79                         return err;
  80         }
  81         return 0;
  82 }
  83
  84 static PyObject *
  85 cycle_next(cycleobject *lz)
  86 {
  87         PyObject *item;
  88         PyObject *it;
  89
  90         while (1) {
  91                 item = PyIter_Next(lz->it);
  92                 if (item != NULL) {
  93                         if (!lz->firstpass)
  94                                 PyList_Append(lz->saved, item);
  95                         return item;
  96                 }
  97                 if (PyErr_Occurred()) {
  98                         if (PyErr_ExceptionMatches(PyExc_StopIteration))
  99                                 PyErr_Clear();
 100                         else
 101                                 return NULL;
 102                 }
 103                 if (PyList_Size(lz->saved) == 0)
 104                         return NULL;
 105                 it = PyObject_GetIter(lz->saved);
 106                 if (it == NULL)
 107                         return NULL;
 108                 Py_DECREF(lz->it);
 109                 lz->it = it;
 110                 lz->firstpass = 1;
 111         }
 112 }
 113
 114 PyDoc_STRVAR(cycle_doc,
 115 "cycle(iterable) --> cycle object\n\
 116 \n\
 117 Return elements from the iterable until it is exhausted.\n\
 118 Then repeat the sequence indefinitely.");
 119
 120 static PyTypeObject cycle_type = {
 121         PyObject_HEAD_INIT(NULL)
 122         0,                              /* ob_size */
 123         "itertools.cycle",              /* tp_name */
 124         sizeof(cycleobject),            /* tp_basicsize */
 125         0,                              /* tp_itemsize */
 126         /* methods */
 127         (destructor)cycle_dealloc,      /* tp_dealloc */
 128         0,                              /* tp_print */
 129         0,                              /* tp_getattr */
 130         0,                              /* tp_setattr */
 131         0,                              /* tp_compare */
 132         0,                              /* tp_repr */
 133         0,                              /* tp_as_number */
 134         0,                              /* tp_as_sequence */
 135         0,                              /* tp_as_mapping */
 136         0,                              /* tp_hash */
 137         0,                              /* tp_call */
 138         0,                              /* tp_str */
 139         PyObject_GenericGetAttr,        /* tp_getattro */
 140         0,                              /* tp_setattro */
 141         0,                              /* tp_as_buffer */
 142         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 143                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 144         cycle_doc,                      /* tp_doc */
 145         (traverseproc)cycle_traverse,   /* tp_traverse */
 146         0,                              /* tp_clear */
 147         0,                              /* tp_richcompare */
 148         0,                              /* tp_weaklistoffset */
 149         PyObject_SelfIter,              /* tp_iter */
 150         (iternextfunc)cycle_next,       /* tp_iternext */
 151         0,                              /* tp_methods */
 152         0,                              /* tp_members */
 153         0,                              /* tp_getset */
 154         0,                              /* tp_base */
 155         0,                              /* tp_dict */
 156         0,                              /* tp_descr_get */
 157         0,                              /* tp_descr_set */
 158         0,                              /* tp_dictoffset */
 159         0,                              /* tp_init */
 160         0,                              /* tp_alloc */
 161         cycle_new,                      /* tp_new */
 162         PyObject_GC_Del,                /* tp_free */
 163 };
 164
 165
 166 /* dropwhile object **********************************************************/
 167
 168 typedef struct {
 169         PyObject_HEAD
 170         PyObject *func;
 171         PyObject *it;
 172         long     start;
 173 } dropwhileobject;
 174
 175 static PyTypeObject dropwhile_type;
 176
 177 static PyObject *
 178 dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 179 {
 180         PyObject *func, *seq;
 181         PyObject *it;
 182         dropwhileobject *lz;
 183
 184         if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
 185                 return NULL;
 186
 187         /* Get iterator. */
 188         it = PyObject_GetIter(seq);
 189         if (it == NULL)
 190                 return NULL;
 191
 192         /* create dropwhileobject structure */
 193         lz = (dropwhileobject *)type->tp_alloc(type, 0);
 194         if (lz == NULL) {
 195                 Py_DECREF(it);
 196                 return NULL;
 197         }
 198         Py_INCREF(func);
 199         lz->func = func;
 200         lz->it = it;
 201         lz->start = 0;
 202
 203         return (PyObject *)lz;
 204 }
 205
 206 static void
 207 dropwhile_dealloc(dropwhileobject *lz)
 208 {
 209         PyObject_GC_UnTrack(lz);
 210         Py_XDECREF(lz->func);
 211         Py_XDECREF(lz->it);
 212         lz->ob_type->tp_free(lz);
 213 }
 214
 215 static int
 216 dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
 217 {
 218         int err;
 219
 220         if (lz->it) {
 221                 err = visit(lz->it, arg);
 222                 if (err)
 223                         return err;
 224         }
 225         if (lz->func) {
 226                 err = visit(lz->func, arg);
 227                 if (err)
 228                         return err;
 229         }
 230         return 0;
 231 }
 232
 233 static PyObject *
 234 dropwhile_next(dropwhileobject *lz)
 235 {
 236         PyObject *item, *good;
 237         PyObject *it = lz->it;
 238         long ok;
 239
 240         for (;;) {
 241                 assert(PyIter_Check(it));
 242                 item = (*it->ob_type->tp_iternext)(it);
 243                 if (item == NULL)
 244                         return NULL;
 245                 if (lz->start == 1)
 246                         return item;
 247
 248                 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
 249                 if (good == NULL) {
 250                         Py_DECREF(item);
 251                         return NULL;
 252                 }
 253                 ok = PyObject_IsTrue(good);
 254                 Py_DECREF(good);
 255                 if (!ok) {
 256                         lz->start = 1;
 257                         return item;
 258                 }
 259                 Py_DECREF(item);
 260         }
 261 }
 262
 263 PyDoc_STRVAR(dropwhile_doc,
 264 "dropwhile(predicate, iterable) --> dropwhile object\n\
 265 \n\
 266 Drop items from the iterable while predicate(item) is true.\n\
 267 Afterwards, return every element until the iterable is exhausted.");
 268
 269 static PyTypeObject dropwhile_type = {
 270         PyObject_HEAD_INIT(NULL)
 271         0,                              /* ob_size */
 272         "itertools.dropwhile",          /* tp_name */
 273         sizeof(dropwhileobject),        /* tp_basicsize */
 274         0,                              /* tp_itemsize */
 275         /* methods */
 276         (destructor)dropwhile_dealloc,  /* tp_dealloc */
 277         0,                              /* tp_print */
 278         0,                              /* tp_getattr */
 279         0,                              /* tp_setattr */
 280         0,                              /* tp_compare */
 281         0,                              /* tp_repr */
 282         0,                              /* tp_as_number */
 283         0,                              /* tp_as_sequence */
 284         0,                              /* tp_as_mapping */
 285         0,                              /* tp_hash */
 286         0,                              /* tp_call */
 287         0,                              /* tp_str */
 288         PyObject_GenericGetAttr,        /* tp_getattro */
 289         0,                              /* tp_setattro */
 290         0,                              /* tp_as_buffer */
 291         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 292                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 293         dropwhile_doc,                  /* tp_doc */
 294         (traverseproc)dropwhile_traverse,    /* tp_traverse */
 295         0,                              /* tp_clear */
 296         0,                              /* tp_richcompare */
 297         0,                              /* tp_weaklistoffset */
 298         PyObject_SelfIter,              /* tp_iter */
 299         (iternextfunc)dropwhile_next,   /* tp_iternext */
 300         0,                              /* tp_methods */
 301         0,                              /* tp_members */
 302         0,                              /* tp_getset */
 303         0,                              /* tp_base */
 304         0,                              /* tp_dict */
 305         0,                              /* tp_descr_get */
 306         0,                              /* tp_descr_set */
 307         0,                              /* tp_dictoffset */
 308         0,                              /* tp_init */
 309         0,                              /* tp_alloc */
 310         dropwhile_new,                  /* tp_new */
 311         PyObject_GC_Del,                /* tp_free */
 312 };
 313
 314
 315 /* takewhile object **********************************************************/
 316
 317 typedef struct {
 318         PyObject_HEAD
 319         PyObject *func;
 320         PyObject *it;
 321         long     stop;
 322 } takewhileobject;
 323
 324 static PyTypeObject takewhile_type;
 325
 326 static PyObject *
 327 takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 328 {
 329         PyObject *func, *seq;
 330         PyObject *it;
 331         takewhileobject *lz;
 332
 333         if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
 334                 return NULL;
 335
 336         /* Get iterator. */
 337         it = PyObject_GetIter(seq);
 338         if (it == NULL)
 339                 return NULL;
 340
 341         /* create takewhileobject structure */
 342         lz = (takewhileobject *)type->tp_alloc(type, 0);
 343         if (lz == NULL) {
 344                 Py_DECREF(it);
 345                 return NULL;
 346         }
 347         Py_INCREF(func);
 348         lz->func = func;
 349         lz->it = it;
 350         lz->stop = 0;
 351
 352         return (PyObject *)lz;
 353 }
 354
 355 static void
 356 takewhile_dealloc(takewhileobject *lz)
 357 {
 358         PyObject_GC_UnTrack(lz);
 359         Py_XDECREF(lz->func);
 360         Py_XDECREF(lz->it);
 361         lz->ob_type->tp_free(lz);
 362 }
 363
 364 static int
 365 takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
 366 {
 367         int err;
 368
 369         if (lz->it) {
 370                 err = visit(lz->it, arg);
 371                 if (err)
 372                         return err;
 373         }
 374         if (lz->func) {
 375                 err = visit(lz->func, arg);
 376                 if (err)
 377                         return err;
 378         }
 379         return 0;
 380 }
 381
 382 static PyObject *
 383 takewhile_next(takewhileobject *lz)
 384 {
 385         PyObject *item, *good;
 386         PyObject *it = lz->it;
 387         long ok;
 388
 389         if (lz->stop == 1)
 390                 return NULL;
 391
 392         assert(PyIter_Check(it));
 393         item = (*it->ob_type->tp_iternext)(it);
 394         if (item == NULL)
 395                 return NULL;
 396
 397         good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
 398         if (good == NULL) {
 399                 Py_DECREF(item);
 400                 return NULL;
 401         }
 402         ok = PyObject_IsTrue(good);
 403         Py_DECREF(good);
 404         if (ok)
 405                 return item;
 406         Py_DECREF(item);
 407         lz->stop = 1;
 408         return NULL;
 409 }
 410
 411 PyDoc_STRVAR(takewhile_doc,
 412 "takewhile(predicate, iterable) --> takewhile object\n\
 413 \n\
 414 Return successive entries from an iterable as long as the \n\
 415 predicate evaluates to true for each entry.");
 416
 417 static PyTypeObject takewhile_type = {
 418         PyObject_HEAD_INIT(NULL)
 419         0,                              /* ob_size */
 420         "itertools.takewhile",          /* tp_name */
 421         sizeof(takewhileobject),        /* tp_basicsize */
 422         0,                              /* tp_itemsize */
 423         /* methods */
 424         (destructor)takewhile_dealloc,  /* tp_dealloc */
 425         0,                              /* tp_print */
 426         0,                              /* tp_getattr */
 427         0,                              /* tp_setattr */
 428         0,                              /* tp_compare */
 429         0,                              /* tp_repr */
 430         0,                              /* tp_as_number */
 431         0,                              /* tp_as_sequence */
 432         0,                              /* tp_as_mapping */
 433         0,                              /* tp_hash */
 434         0,                              /* tp_call */
 435         0,                              /* tp_str */
 436         PyObject_GenericGetAttr,        /* tp_getattro */
 437         0,                              /* tp_setattro */
 438         0,                              /* tp_as_buffer */
 439         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 440                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 441         takewhile_doc,                  /* tp_doc */
 442         (traverseproc)takewhile_traverse,    /* tp_traverse */
 443         0,                              /* tp_clear */
 444         0,                              /* tp_richcompare */
 445         0,                              /* tp_weaklistoffset */
 446         PyObject_SelfIter,              /* tp_iter */
 447         (iternextfunc)takewhile_next,   /* tp_iternext */
 448         0,                              /* tp_methods */
 449         0,                              /* tp_members */
 450         0,                              /* tp_getset */
 451         0,                              /* tp_base */
 452         0,                              /* tp_dict */
 453         0,                              /* tp_descr_get */
 454         0,                              /* tp_descr_set */
 455         0,                              /* tp_dictoffset */
 456         0,                              /* tp_init */
 457         0,                              /* tp_alloc */
 458         takewhile_new,                  /* tp_new */
 459         PyObject_GC_Del,                /* tp_free */
 460 };
 461
 462
 463 /* islice object ************************************************************/
 464
 465 typedef struct {
 466         PyObject_HEAD
 467         PyObject *it;
 468         long    next;
 469         long    stop;
 470         long    step;
 471         long    cnt;
 472 } isliceobject;
 473
 474 static PyTypeObject islice_type;
 475
 476 static PyObject *
 477 islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 478 {
 479         PyObject *seq;
 480         long start=0, stop=-1, step=1;
 481         PyObject *it, *a1=NULL, *a2=NULL;
 482         int numargs;
 483         isliceobject *lz;
 484
 485         numargs = PyTuple_Size(args);
 486         if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
 487                 return NULL;
 488
 489         if (numargs == 2) {
 490                 if (a1 != Py_None) {
 491                         stop = PyInt_AsLong(a1);
 492                         if (stop == -1) {
 493                                 if (PyErr_Occurred())
 494                                         PyErr_Clear();
 495                                 PyErr_SetString(PyExc_ValueError,
 496                                    "Stop argument must be an integer or None.");
 497                                 return NULL;
 498                         }
 499                 }
 500         } else {
 501                 start = PyInt_AsLong(a1);
 502                 if (start == -1 && PyErr_Occurred()) {
 503                         PyErr_Clear();
 504                         PyErr_SetString(PyExc_ValueError,
 505                            "Start argument must be an integer.");
 506                         return NULL;
 507                 }
 508                 if (a2 != Py_None) {
 509                         stop = PyInt_AsLong(a2);
 510                         if (stop == -1) {
 511                                 if (PyErr_Occurred())
 512                                         PyErr_Clear();
 513                                 PyErr_SetString(PyExc_ValueError,
 514                                    "Stop argument must be an integer or None.");
 515                                 return NULL;
 516                         }
 517                 }
 518         }
 519
 520         if (start<0 || stop<-1) {
 521                 PyErr_SetString(PyExc_ValueError,
 522                    "Indices for islice() must be positive.");
 523                 return NULL;
 524         }
 525
 526         if (step<1) {
 527                 PyErr_SetString(PyExc_ValueError,
 528                    "Step must be one or larger for islice().");
 529                 return NULL;
 530         }
 531
 532         /* Get iterator. */
 533         it = PyObject_GetIter(seq);
 534         if (it == NULL)
 535                 return NULL;
 536
 537         /* create isliceobject structure */
 538         lz = (isliceobject *)type->tp_alloc(type, 0);
 539         if (lz == NULL) {
 540                 Py_DECREF(it);
 541                 return NULL;
 542         }
 543         lz->it = it;
 544         lz->next = start;
 545         lz->stop = stop;
 546         lz->step = step;
 547         lz->cnt = 0L;
 548
 549         return (PyObject *)lz;
 550 }
 551
 552 static void
 553 islice_dealloc(isliceobject *lz)
 554 {
 555         PyObject_GC_UnTrack(lz);
 556         Py_XDECREF(lz->it);
 557         lz->ob_type->tp_free(lz);
 558 }
 559
 560 static int
 561 islice_traverse(isliceobject *lz, visitproc visit, void *arg)
 562 {
 563         if (lz->it)
 564                 return visit(lz->it, arg);
 565         return 0;
 566 }
 567
 568 static PyObject *
 569 islice_next(isliceobject *lz)
 570 {
 571         PyObject *item;
 572         PyObject *it = lz->it;
 573         long oldnext;
 574
 575         while (lz->cnt < lz->next) {
 576                 assert(PyIter_Check(it));
 577                 item = (*it->ob_type->tp_iternext)(it);
 578                 if (item == NULL)
 579                         return NULL;
 580                 Py_DECREF(item);
 581                 lz->cnt++;
 582         }
 583         if (lz->stop != -1 && lz->cnt >= lz->stop)
 584                 return NULL;
 585         assert(PyIter_Check(it));
 586         item = (*it->ob_type->tp_iternext)(it);
 587         if (item == NULL)
 588                 return NULL;
 589         lz->cnt++;
 590         oldnext = lz->next;
 591         lz->next += lz->step;
 592         if (lz->next < oldnext) /* Check for overflow */
 593                 lz->next = lz->stop;
 594         return item;
 595 }
 596
 597 PyDoc_STRVAR(islice_doc,
 598 "islice(iterable, [start,] stop [, step]) --> islice object\n\
 599 \n\
 600 Return an iterator whose next() method returns selected values from an\n\
 601 iterable.  If start is specified, will skip all preceding elements;\n\
 602 otherwise, start defaults to zero.  Step defaults to one.  If\n\
 603 specified as another value, step determines how many values are \n\
 604 skipped between successive calls.  Works like a slice() on a list\n\
 605 but returns an iterator.");
 606
 607 static PyTypeObject islice_type = {
 608         PyObject_HEAD_INIT(NULL)
 609         0,                              /* ob_size */
 610         "itertools.islice",             /* tp_name */
 611         sizeof(isliceobject),           /* tp_basicsize */
 612         0,                              /* tp_itemsize */
 613         /* methods */
 614         (destructor)islice_dealloc,     /* tp_dealloc */
 615         0,                              /* tp_print */
 616         0,                              /* tp_getattr */
 617         0,                              /* tp_setattr */
 618         0,                              /* tp_compare */
 619         0,                              /* tp_repr */
 620         0,                              /* tp_as_number */
 621         0,                              /* tp_as_sequence */
 622         0,                              /* tp_as_mapping */
 623         0,                              /* tp_hash */
 624         0,                              /* tp_call */
 625         0,                              /* tp_str */
 626         PyObject_GenericGetAttr,        /* tp_getattro */
 627         0,                              /* tp_setattro */
 628         0,                              /* tp_as_buffer */
 629         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 630                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 631         islice_doc,                     /* tp_doc */
 632         (traverseproc)islice_traverse,  /* tp_traverse */
 633         0,                              /* tp_clear */
 634         0,                              /* tp_richcompare */
 635         0,                              /* tp_weaklistoffset */
 636         PyObject_SelfIter,              /* tp_iter */
 637         (iternextfunc)islice_next,      /* tp_iternext */
 638         0,                              /* tp_methods */
 639         0,                              /* tp_members */
 640         0,                              /* tp_getset */
 641         0,                              /* tp_base */
 642         0,                              /* tp_dict */
 643         0,                              /* tp_descr_get */
 644         0,                              /* tp_descr_set */
 645         0,                              /* tp_dictoffset */
 646         0,                              /* tp_init */
 647         0,                              /* tp_alloc */
 648         islice_new,                     /* tp_new */
 649         PyObject_GC_Del,                /* tp_free */
 650 };
 651
 652
 653 /* starmap object ************************************************************/
 654
 655 typedef struct {
 656         PyObject_HEAD
 657         PyObject *func;
 658         PyObject *it;
 659 } starmapobject;
 660
 661 static PyTypeObject starmap_type;
 662
 663 static PyObject *
 664 starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 665 {
 666         PyObject *func, *seq;
 667         PyObject *it;
 668         starmapobject *lz;
 669
 670         if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
 671                 return NULL;
 672
 673         /* Get iterator. */
 674         it = PyObject_GetIter(seq);
 675         if (it == NULL)
 676                 return NULL;
 677
 678         /* create starmapobject structure */
 679         lz = (starmapobject *)type->tp_alloc(type, 0);
 680         if (lz == NULL) {
 681                 Py_DECREF(it);
 682                 return NULL;
 683         }
 684         Py_INCREF(func);
 685         lz->func = func;
 686         lz->it = it;
 687
 688         return (PyObject *)lz;
 689 }
 690
 691 static void
 692 starmap_dealloc(starmapobject *lz)
 693 {
 694         PyObject_GC_UnTrack(lz);
 695         Py_XDECREF(lz->func);
 696         Py_XDECREF(lz->it);
 697         lz->ob_type->tp_free(lz);
 698 }
 699
 700 static int
 701 starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
 702 {
 703         int err;
 704
 705         if (lz->it) {
 706                 err = visit(lz->it, arg);
 707                 if (err)
 708                         return err;
 709         }
 710         if (lz->func) {
 711                 err = visit(lz->func, arg);
 712                 if (err)
 713                         return err;
 714         }
 715         return 0;
 716 }
 717
 718 static PyObject *
 719 starmap_next(starmapobject *lz)
 720 {
 721         PyObject *args;
 722         PyObject *result;
 723         PyObject *it = lz->it;
 724
 725         assert(PyIter_Check(it));
 726         args = (*it->ob_type->tp_iternext)(it);
 727         if (args == NULL)
 728                 return NULL;
 729         if (!PyTuple_CheckExact(args)) {
 730                 Py_DECREF(args);
 731                 PyErr_SetString(PyExc_TypeError,
 732                                 "iterator must return a tuple");
 733                 return NULL;
 734         }
 735         result = PyObject_Call(lz->func, args, NULL);
 736         Py_DECREF(args);
 737         return result;
 738 }
 739
 740 PyDoc_STRVAR(starmap_doc,
 741 "starmap(function, sequence) --> starmap object\n\
 742 \n\
 743 Return an iterator whose values are returned from the function evaluated\n\
 744 with a argument tuple taken from the given sequence.");
 745
 746 static PyTypeObject starmap_type = {
 747         PyObject_HEAD_INIT(NULL)
 748         0,                              /* ob_size */
 749         "itertools.starmap",            /* tp_name */
 750         sizeof(starmapobject),          /* tp_basicsize */
 751         0,                              /* tp_itemsize */
 752         /* methods */
 753         (destructor)starmap_dealloc,    /* tp_dealloc */
 754         0,                              /* tp_print */
 755         0,                              /* tp_getattr */
 756         0,                              /* tp_setattr */
 757         0,                              /* tp_compare */
 758         0,                              /* tp_repr */
 759         0,                              /* tp_as_number */
 760         0,                              /* tp_as_sequence */
 761         0,                              /* tp_as_mapping */
 762         0,                              /* tp_hash */
 763         0,                              /* tp_call */
 764         0,                              /* tp_str */
 765         PyObject_GenericGetAttr,        /* tp_getattro */
 766         0,                              /* tp_setattro */
 767         0,                              /* tp_as_buffer */
 768         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 769                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 770         starmap_doc,                    /* tp_doc */
 771         (traverseproc)starmap_traverse, /* tp_traverse */
 772         0,                              /* tp_clear */
 773         0,                              /* tp_richcompare */
 774         0,                              /* tp_weaklistoffset */
 775         PyObject_SelfIter,              /* tp_iter */
 776         (iternextfunc)starmap_next,     /* tp_iternext */
 777         0,                              /* tp_methods */
 778         0,                              /* tp_members */
 779         0,                              /* tp_getset */
 780         0,                              /* tp_base */
 781         0,                              /* tp_dict */
 782         0,                              /* tp_descr_get */
 783         0,                              /* tp_descr_set */
 784         0,                              /* tp_dictoffset */
 785         0,                              /* tp_init */
 786         0,                              /* tp_alloc */
 787         starmap_new,                    /* tp_new */
 788         PyObject_GC_Del,                /* tp_free */
 789 };
 790
 791
 792 /* imap object ************************************************************/
 793
 794 typedef struct {
 795         PyObject_HEAD
 796         PyObject *iters;
 797         PyObject *func;
 798 } imapobject;
 799
 800 static PyTypeObject imap_type;
 801
 802 static PyObject *
 803 imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 804 {
 805         PyObject *it, *iters, *func;
 806         imapobject *lz;
 807         int numargs, i;
 808
 809         numargs = PyTuple_Size(args);
 810         if (numargs < 2) {
 811                 PyErr_SetString(PyExc_TypeError,
 812                    "imap() must have at least two arguments.");
 813                 return NULL;
 814         }
 815
 816         iters = PyTuple_New(numargs-1);
 817         if (iters == NULL)
 818                 return NULL;
 819
 820         for (i=1 ; i<numargs ; i++) {
 821                 /* Get iterator. */
 822                 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
 823                 if (it == NULL) {
 824                         Py_DECREF(iters);
 825                         return NULL;
 826                 }
 827                 PyTuple_SET_ITEM(iters, i-1, it);
 828         }
 829
 830         /* create imapobject structure */
 831         lz = (imapobject *)type->tp_alloc(type, 0);
 832         if (lz == NULL) {
 833                 Py_DECREF(iters);
 834                 return NULL;
 835         }
 836         lz->iters = iters;
 837         func = PyTuple_GET_ITEM(args, 0);
 838         Py_INCREF(func);
 839         lz->func = func;
 840
 841         return (PyObject *)lz;
 842 }
 843
 844 static void
 845 imap_dealloc(imapobject *lz)
 846 {
 847         PyObject_GC_UnTrack(lz);
 848         Py_XDECREF(lz->iters);
 849         Py_XDECREF(lz->func);
 850         lz->ob_type->tp_free(lz);
 851 }
 852
 853 static int
 854 imap_traverse(imapobject *lz, visitproc visit, void *arg)
 855 {
 856         int err;
 857
 858         if (lz->iters) {
 859                 err = visit(lz->iters, arg);
 860                 if (err)
 861                         return err;
 862         }
 863         if (lz->func) {
 864                 err = visit(lz->func, arg);
 865                 if (err)
 866                         return err;
 867         }
 868         return 0;
 869 }
 870
 871 /*
 872 imap() is an iterator version of __builtins__.map() except that it does
 873 not have the None fill-in feature.  That was intentionally left out for
 874 the following reasons:
 875
 876   1) Itertools are designed to be easily combined and chained together.
 877      Having all tools stop with the shortest input is a unifying principle
 878      that makes it easier to combine finite iterators (supplying data) with
 879      infinite iterators like count() and repeat() (for supplying sequential
 880      or constant arguments to a function).
 881
 882   2) In typical use cases for combining itertools, having one finite data
 883      supplier run out before another is likely to be an error condition which
 884      should not pass silently by automatically supplying None.
 885
 886   3) The use cases for automatic None fill-in are rare -- not many functions
 887      do something useful when a parameter suddenly switches type and becomes
 888      None.
 889
 890   4) If a need does arise, it can be met by __builtins__.map() or by
 891      writing:  chain(iterable, repeat(None)).
 892
 893   5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
 894 */
 895
 896 static PyObject *
 897 imap_next(imapobject *lz)
 898 {
 899         PyObject *val;
 900         PyObject *argtuple;
 901         PyObject *result;
 902         int numargs, i;
 903
 904         numargs = PyTuple_Size(lz->iters);
 905         argtuple = PyTuple_New(numargs);
 906         if (argtuple == NULL)
 907                 return NULL;
 908
 909         for (i=0 ; i<numargs ; i++) {
 910                 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
 911                 if (val == NULL) {
 912                         Py_DECREF(argtuple);
 913                         return NULL;
 914                 }
 915                 PyTuple_SET_ITEM(argtuple, i, val);
 916         }
 917         if (lz->func == Py_None)
 918                 return argtuple;
 919         result = PyObject_Call(lz->func, argtuple, NULL);
 920         Py_DECREF(argtuple);
 921         return result;
 922 }
 923
 924 PyDoc_STRVAR(imap_doc,
 925 "imap(func, *iterables) --> imap object\n\
 926 \n\
 927 Make an iterator that computes the function using arguments from\n\
 928 each of the iterables.  Like map() except that it returns\n\
 929 an iterator instead of a list and that it stops when the shortest\n\
 930 iterable is exhausted instead of filling in None for shorter\n\
 931 iterables.");
 932
 933 static PyTypeObject imap_type = {
 934         PyObject_HEAD_INIT(NULL)
 935         0,                              /* ob_size */
 936         "itertools.imap",               /* tp_name */
 937         sizeof(imapobject),             /* tp_basicsize */
 938         0,                              /* tp_itemsize */
 939         /* methods */
 940         (destructor)imap_dealloc,       /* tp_dealloc */
 941         0,                              /* tp_print */
 942         0,                              /* tp_getattr */
 943         0,                              /* tp_setattr */
 944         0,                              /* tp_compare */
 945         0,                              /* tp_repr */
 946         0,                              /* tp_as_number */
 947         0,                              /* tp_as_sequence */
 948         0,                              /* tp_as_mapping */
 949         0,                              /* tp_hash */
 950         0,                              /* tp_call */
 951         0,                              /* tp_str */
 952         PyObject_GenericGetAttr,        /* tp_getattro */
 953         0,                              /* tp_setattro */
 954         0,                              /* tp_as_buffer */
 955         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 956                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 957         imap_doc,                       /* tp_doc */
 958         (traverseproc)imap_traverse,    /* tp_traverse */
 959         0,                              /* tp_clear */
 960         0,                              /* tp_richcompare */
 961         0,                              /* tp_weaklistoffset */
 962         PyObject_SelfIter,              /* tp_iter */
 963         (iternextfunc)imap_next,        /* tp_iternext */
 964         0,                              /* tp_methods */
 965         0,                              /* tp_members */
 966         0,                              /* tp_getset */
 967         0,                              /* tp_base */
 968         0,                              /* tp_dict */
 969         0,                              /* tp_descr_get */
 970         0,                              /* tp_descr_set */
 971         0,                              /* tp_dictoffset */
 972         0,                              /* tp_init */
 973         0,                              /* tp_alloc */
 974         imap_new,                       /* tp_new */
 975         PyObject_GC_Del,                /* tp_free */
 976 };
 977
 978
 979 /* chain object ************************************************************/
 980
 981 typedef struct {
 982         PyObject_HEAD
 983         long    tuplesize;
 984         long    iternum;                /* which iterator is active */
 985         PyObject *ittuple;              /* tuple of iterators */
 986 } chainobject;
 987
 988 static PyTypeObject chain_type;
 989
 990 static PyObject *
 991 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 992 {
 993         chainobject *lz;
 994         int tuplesize = PySequence_Length(args);
 995         int i;
 996         PyObject *ittuple;
 997
 998         /* obtain iterators */
 999         assert(PyTuple_Check(args));
1000         ittuple = PyTuple_New(tuplesize);
1001         if(ittuple == NULL)
1002                 return NULL;
1003         for (i=0; i < tuplesize; ++i) {
1004                 PyObject *item = PyTuple_GET_ITEM(args, i);
1005                 PyObject *it = PyObject_GetIter(item);
1006                 if (it == NULL) {
1007                         if (PyErr_ExceptionMatches(PyExc_TypeError))
1008                                 PyErr_Format(PyExc_TypeError,
1009                                     "chain argument #%d must support iteration",
1010                                     i+1);
1011                         Py_DECREF(ittuple);
1012                         return NULL;
1013                 }
1014                 PyTuple_SET_ITEM(ittuple, i, it);
1015         }
1016
1017         /* create chainobject structure */
1018         lz = (chainobject *)type->tp_alloc(type, 0);
1019         if (lz == NULL) {
1020                 Py_DECREF(ittuple);
1021                 return NULL;
1022         }
1023
1024         lz->ittuple = ittuple;
1025         lz->iternum = 0;
1026         lz->tuplesize = tuplesize;
1027
1028         return (PyObject *)lz;
1029 }
1030
1031 static void
1032 chain_dealloc(chainobject *lz)
1033 {
1034         PyObject_GC_UnTrack(lz);
1035         Py_XDECREF(lz->ittuple);
1036         lz->ob_type->tp_free(lz);
1037 }
1038
1039 static int
1040 chain_traverse(chainobject *lz, visitproc visit, void *arg)
1041 {
1042         if (lz->ittuple)
1043                 return visit(lz->ittuple, arg);
1044         return 0;
1045 }
1046
1047 static PyObject *
1048 chain_next(chainobject *lz)
1049 {
1050         PyObject *it;
1051         PyObject *item;
1052
1053         while (lz->iternum < lz->tuplesize) {
1054                 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1055                 item = PyIter_Next(it);
1056                 if (item != NULL)
1057                         return item;
1058                 if (PyErr_Occurred()) {
1059                         if (PyErr_ExceptionMatches(PyExc_StopIteration))
1060                                 PyErr_Clear();
1061                         else
1062                                 return NULL;
1063                 }
1064                 lz->iternum++;
1065         }
1066         return NULL;
1067 }
1068
1069 PyDoc_STRVAR(chain_doc,
1070 "chain(*iterables) --> chain object\n\
1071 \n\
1072 Return a chain object whose .next() method returns elements from the\n\
1073 first iterable until it is exhausted, then elements from the next\n\
1074 iterable, until all of the iterables are exhausted.");
1075
1076 static PyTypeObject chain_type = {
1077         PyObject_HEAD_INIT(NULL)
1078         0,                              /* ob_size */
1079         "itertools.chain",              /* tp_name */
1080         sizeof(chainobject),            /* tp_basicsize */
1081         0,                              /* tp_itemsize */
1082         /* methods */
1083         (destructor)chain_dealloc,      /* tp_dealloc */
1084         0,                              /* tp_print */
1085         0,                              /* tp_getattr */
1086         0,                              /* tp_setattr */
1087         0,                              /* tp_compare */
1088         0,                              /* tp_repr */
1089         0,                              /* tp_as_number */
1090         0,                              /* tp_as_sequence */
1091         0,                              /* tp_as_mapping */
1092         0,                              /* tp_hash */
1093         0,                              /* tp_call */
1094         0,                              /* tp_str */
1095         PyObject_GenericGetAttr,        /* tp_getattro */
1096         0,                              /* tp_setattro */
1097         0,                              /* tp_as_buffer */
1098         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1099                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1100         chain_doc,                      /* tp_doc */
1101         (traverseproc)chain_traverse,   /* tp_traverse */
1102         0,                              /* tp_clear */
1103         0,                              /* tp_richcompare */
1104         0,                              /* tp_weaklistoffset */
1105         PyObject_SelfIter,              /* tp_iter */
1106         (iternextfunc)chain_next,       /* tp_iternext */
1107         0,                              /* tp_methods */
1108         0,                              /* tp_members */
1109         0,                              /* tp_getset */
1110         0,                              /* tp_base */
1111         0,                              /* tp_dict */
1112         0,                              /* tp_descr_get */
1113         0,                              /* tp_descr_set */
1114         0,                              /* tp_dictoffset */
1115         0,                              /* tp_init */
1116         0,                              /* tp_alloc */
1117         chain_new,                      /* tp_new */
1118         PyObject_GC_Del,                /* tp_free */
1119 };
1120
1121
1122 /* ifilter object ************************************************************/
1123
1124 typedef struct {
1125         PyObject_HEAD
1126         PyObject *func;
1127         PyObject *it;
1128 } ifilterobject;
1129
1130 static PyTypeObject ifilter_type;
1131
1132 static PyObject *
1133 ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1134 {
1135         PyObject *func, *seq;
1136         PyObject *it;
1137         ifilterobject *lz;
1138
1139         if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
1140                 return NULL;
1141
1142         /* Get iterator. */
1143         it = PyObject_GetIter(seq);
1144         if (it == NULL)
1145                 return NULL;
1146
1147         /* create ifilterobject structure */
1148         lz = (ifilterobject *)type->tp_alloc(type, 0);
1149         if (lz == NULL) {
1150                 Py_DECREF(it);
1151                 return NULL;
1152         }
1153         Py_INCREF(func);
1154         lz->func = func;
1155         lz->it = it;
1156
1157         return (PyObject *)lz;
1158 }
1159
1160 static void
1161 ifilter_dealloc(ifilterobject *lz)
1162 {
1163         PyObject_GC_UnTrack(lz);
1164         Py_XDECREF(lz->func);
1165         Py_XDECREF(lz->it);
1166         lz->ob_type->tp_free(lz);
1167 }
1168
1169 static int
1170 ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1171 {
1172         int err;
1173
1174         if (lz->it) {
1175                 err = visit(lz->it, arg);
1176                 if (err)
1177                         return err;
1178         }
1179         if (lz->func) {
1180                 err = visit(lz->func, arg);
1181                 if (err)
1182                         return err;
1183         }
1184         return 0;
1185 }
1186
1187 static PyObject *
1188 ifilter_next(ifilterobject *lz)
1189 {
1190         PyObject *item;
1191         PyObject *it = lz->it;
1192         long ok;
1193
1194         for (;;) {
1195                 assert(PyIter_Check(it));
1196                 item = (*it->ob_type->tp_iternext)(it);
1197                 if (item == NULL)
1198                         return NULL;
1199
1200                 if (lz->func == Py_None) {
1201                         ok = PyObject_IsTrue(item);
1202                 } else {
1203                         PyObject *good;
1204                         good = PyObject_CallFunctionObjArgs(lz->func,
1205                                                             item, NULL);
1206                         if (good == NULL) {
1207                                 Py_DECREF(item);
1208                                 return NULL;
1209                         }
1210                         ok = PyObject_IsTrue(good);
1211                         Py_DECREF(good);
1212                 }
1213                 if (ok)
1214                         return item;
1215                 Py_DECREF(item);
1216         }
1217 }
1218
1219 PyDoc_STRVAR(ifilter_doc,
1220 "ifilter(function or None, sequence) --> ifilter object\n\
1221 \n\
1222 Return those items of sequence for which function(item) is true.\n\
1223 If function is None, return the items that are true.");
1224
1225 static PyTypeObject ifilter_type = {
1226         PyObject_HEAD_INIT(NULL)
1227         0,                              /* ob_size */
1228         "itertools.ifilter",            /* tp_name */
1229         sizeof(ifilterobject),          /* tp_basicsize */
1230         0,                              /* tp_itemsize */
1231         /* methods */
1232         (destructor)ifilter_dealloc,    /* tp_dealloc */
1233         0,                              /* tp_print */
1234         0,                              /* tp_getattr */
1235         0,                              /* tp_setattr */
1236         0,                              /* tp_compare */
1237         0,                              /* tp_repr */
1238         0,                              /* tp_as_number */
1239         0,                              /* tp_as_sequence */
1240         0,                              /* tp_as_mapping */
1241         0,                              /* tp_hash */
1242         0,                              /* tp_call */
1243         0,                              /* tp_str */
1244         PyObject_GenericGetAttr,        /* tp_getattro */
1245         0,                              /* tp_setattro */
1246         0,                              /* tp_as_buffer */
1247         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1248                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1249         ifilter_doc,                    /* tp_doc */
1250         (traverseproc)ifilter_traverse, /* tp_traverse */
1251         0,                              /* tp_clear */
1252         0,                              /* tp_richcompare */
1253         0,                              /* tp_weaklistoffset */
1254         PyObject_SelfIter,              /* tp_iter */
1255         (iternextfunc)ifilter_next,     /* tp_iternext */
1256         0,                              /* tp_methods */
1257         0,                              /* tp_members */
1258         0,                              /* tp_getset */
1259         0,                              /* tp_base */
1260         0,                              /* tp_dict */
1261         0,                              /* tp_descr_get */
1262         0,                              /* tp_descr_set */
1263         0,                              /* tp_dictoffset */
1264         0,                              /* tp_init */
1265         0,                              /* tp_alloc */
1266         ifilter_new,                    /* tp_new */
1267         PyObject_GC_Del,                /* tp_free */
1268 };
1269
1270
1271 /* ifilterfalse object ************************************************************/
1272
1273 typedef struct {
1274         PyObject_HEAD
1275         PyObject *func;
1276         PyObject *it;
1277 } ifilterfalseobject;
1278
1279 static PyTypeObject ifilterfalse_type;
1280
1281 static PyObject *
1282 ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1283 {
1284         PyObject *func, *seq;
1285         PyObject *it;
1286         ifilterfalseobject *lz;
1287
1288         if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1289                 return NULL;
1290
1291         /* Get iterator. */
1292         it = PyObject_GetIter(seq);
1293         if (it == NULL)
1294                 return NULL;
1295
1296         /* create ifilterfalseobject structure */
1297         lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1298         if (lz == NULL) {
1299                 Py_DECREF(it);
1300                 return NULL;
1301         }
1302         Py_INCREF(func);
1303         lz->func = func;
1304         lz->it = it;
1305
1306         return (PyObject *)lz;
1307 }
1308
1309 static void
1310 ifilterfalse_dealloc(ifilterfalseobject *lz)
1311 {
1312         PyObject_GC_UnTrack(lz);
1313         Py_XDECREF(lz->func);
1314         Py_XDECREF(lz->it);
1315         lz->ob_type->tp_free(lz);
1316 }
1317
1318 static int
1319 ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1320 {
1321         int err;
1322
1323         if (lz->it) {
1324                 err = visit(lz->it, arg);
1325                 if (err)
1326                         return err;
1327         }
1328         if (lz->func) {
1329                 err = visit(lz->func, arg);
1330                 if (err)
1331                         return err;
1332         }
1333         return 0;
1334 }
1335
1336 static PyObject *
1337 ifilterfalse_next(ifilterfalseobject *lz)
1338 {
1339         PyObject *item;
1340         PyObject *it = lz->it;
1341         long ok;
1342
1343         for (;;) {
1344                 assert(PyIter_Check(it));
1345                 item = (*it->ob_type->tp_iternext)(it);
1346                 if (item == NULL)
1347                         return NULL;
1348
1349                 if (lz->func == Py_None) {
1350                         ok = PyObject_IsTrue(item);
1351                 } else {
1352                         PyObject *good;
1353                         good = PyObject_CallFunctionObjArgs(lz->func,
1354                                                             item, NULL);
1355                         if (good == NULL) {
1356                                 Py_DECREF(item);
1357                                 return NULL;
1358                         }
1359                         ok = PyObject_IsTrue(good);
1360                         Py_DECREF(good);
1361                 }
1362                 if (!ok)
1363                         return item;
1364                 Py_DECREF(item);
1365         }
1366 }
1367
1368 PyDoc_STRVAR(ifilterfalse_doc,
1369 "ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1370 \n\
1371 Return those items of sequence for which function(item) is false.\n\
1372 If function is None, return the items that are false.");
1373
1374 static PyTypeObject ifilterfalse_type = {
1375         PyObject_HEAD_INIT(NULL)
1376         0,                              /* ob_size */
1377         "itertools.ifilterfalse",       /* tp_name */
1378         sizeof(ifilterfalseobject),     /* tp_basicsize */
1379         0,                              /* tp_itemsize */
1380         /* methods */
1381         (destructor)ifilterfalse_dealloc,       /* tp_dealloc */
1382         0,                              /* tp_print */
1383         0,                              /* tp_getattr */
1384         0,                              /* tp_setattr */
1385         0,                              /* tp_compare */
1386         0,                              /* tp_repr */
1387         0,                              /* tp_as_number */
1388         0,                              /* tp_as_sequence */
1389         0,                              /* tp_as_mapping */
1390         0,                              /* tp_hash */
1391         0,                              /* tp_call */
1392         0,                              /* tp_str */
1393         PyObject_GenericGetAttr,        /* tp_getattro */
1394         0,                              /* tp_setattro */
1395         0,                              /* tp_as_buffer */
1396         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1397                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1398         ifilterfalse_doc,               /* tp_doc */
1399         (traverseproc)ifilterfalse_traverse,    /* tp_traverse */
1400         0,                              /* tp_clear */
1401         0,                              /* tp_richcompare */
1402         0,                              /* tp_weaklistoffset */
1403         PyObject_SelfIter,              /* tp_iter */
1404         (iternextfunc)ifilterfalse_next,        /* tp_iternext */
1405         0,                              /* tp_methods */
1406         0,                              /* tp_members */
1407         0,                              /* tp_getset */
1408         0,                              /* tp_base */
1409         0,                              /* tp_dict */
1410         0,                              /* tp_descr_get */
1411         0,                              /* tp_descr_set */
1412         0,                              /* tp_dictoffset */
1413         0,                              /* tp_init */
1414         0,                              /* tp_alloc */
1415         ifilterfalse_new,               /* tp_new */
1416         PyObject_GC_Del,                /* tp_free */
1417 };
1418
1419
1420 /* count object ************************************************************/
1421
1422 typedef struct {
1423         PyObject_HEAD
1424         long    cnt;
1425 } countobject;
1426
1427 static PyTypeObject count_type;
1428
1429 static PyObject *
1430 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1431 {
1432         countobject *lz;
1433         long cnt = 0;
1434
1435         if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1436                 return NULL;
1437
1438         /* create countobject structure */
1439         lz = (countobject *)PyObject_New(countobject, &count_type);
1440         if (lz == NULL)
1441                 return NULL;
1442         lz->cnt = cnt;
1443
1444         return (PyObject *)lz;
1445 }
1446
1447 static PyObject *
1448 count_next(countobject *lz)
1449 {
1450         return PyInt_FromLong(lz->cnt++);
1451 }
1452
1453 PyDoc_STRVAR(count_doc,
1454 "count([firstval]) --> count object\n\
1455 \n\
1456 Return a count object whose .next() method returns consecutive\n\
1457 integers starting from zero or, if specified, from firstval.");
1458
1459 static PyTypeObject count_type = {
1460         PyObject_HEAD_INIT(NULL)
1461         0,                              /* ob_size */
1462         "itertools.count",              /* tp_name */
1463         sizeof(countobject),            /* tp_basicsize */
1464         0,                              /* tp_itemsize */
1465         /* methods */
1466         (destructor)PyObject_Del,       /* tp_dealloc */
1467         0,                              /* tp_print */
1468         0,                              /* tp_getattr */
1469         0,                              /* tp_setattr */
1470         0,                              /* tp_compare */
1471         0,                              /* tp_repr */
1472         0,                              /* tp_as_number */
1473         0,                              /* tp_as_sequence */
1474         0,                              /* tp_as_mapping */
1475         0,                              /* tp_hash */
1476         0,                              /* tp_call */
1477         0,                              /* tp_str */
1478         PyObject_GenericGetAttr,        /* tp_getattro */
1479         0,                              /* tp_setattro */
1480         0,                              /* tp_as_buffer */
1481         Py_TPFLAGS_DEFAULT,             /* tp_flags */
1482         count_doc,                      /* tp_doc */
1483         0,                              /* tp_traverse */
1484         0,                              /* tp_clear */
1485         0,                              /* tp_richcompare */
1486         0,                              /* tp_weaklistoffset */
1487         PyObject_SelfIter,              /* tp_iter */
1488         (iternextfunc)count_next,       /* tp_iternext */
1489         0,                              /* tp_methods */
1490         0,                              /* tp_members */
1491         0,                              /* tp_getset */
1492         0,                              /* tp_base */
1493         0,                              /* tp_dict */
1494         0,                              /* tp_descr_get */
1495         0,                              /* tp_descr_set */
1496         0,                              /* tp_dictoffset */
1497         0,                              /* tp_init */
1498         0,                              /* tp_alloc */
1499         count_new,                      /* tp_new */
1500 };
1501
1502
1503 /* izip object ************************************************************/
1504
1505 #include "Python.h"
1506
1507 typedef struct {
1508         PyObject_HEAD
1509         long    tuplesize;
1510         PyObject *ittuple;              /* tuple of iterators */
1511         PyObject *result;
1512 } izipobject;
1513
1514 static PyTypeObject izip_type;
1515
1516 static PyObject *
1517 izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1518 {
1519         izipobject *lz;
1520         int i;
1521         PyObject *ittuple;  /* tuple of iterators */
1522         PyObject *result;
1523         int tuplesize = PySequence_Length(args);
1524
1525         /* args must be a tuple */
1526         assert(PyTuple_Check(args));
1527
1528         /* obtain iterators */
1529         ittuple = PyTuple_New(tuplesize);
1530         if(ittuple == NULL)
1531                 return NULL;
1532         for (i=0; i < tuplesize; ++i) {
1533                 PyObject *item = PyTuple_GET_ITEM(args, i);
1534                 PyObject *it = PyObject_GetIter(item);
1535                 if (it == NULL) {
1536                         if (PyErr_ExceptionMatches(PyExc_TypeError))
1537                                 PyErr_Format(PyExc_TypeError,
1538                                     "izip argument #%d must support iteration",
1539                                     i+1);
1540                         Py_DECREF(ittuple);
1541                         return NULL;
1542                 }
1543                 PyTuple_SET_ITEM(ittuple, i, it);
1544         }
1545
1546         /* create a result holder */
1547         result = PyTuple_New(tuplesize);
1548         if (result == NULL) {
1549                 Py_DECREF(ittuple);
1550                 return NULL;
1551         }
1552         for (i=0 ; i < tuplesize ; i++) {
1553                 Py_INCREF(Py_None);
1554                 PyTuple_SET_ITEM(result, i, Py_None);
1555         }
1556
1557         /* create izipobject structure */
1558         lz = (izipobject *)type->tp_alloc(type, 0);
1559         if (lz == NULL) {
1560                 Py_DECREF(ittuple);
1561                 Py_DECREF(result);
1562                 return NULL;
1563         }
1564         lz->ittuple = ittuple;
1565         lz->tuplesize = tuplesize;
1566         lz->result = result;
1567
1568         return (PyObject *)lz;
1569 }
1570
1571 static void
1572 izip_dealloc(izipobject *lz)
1573 {
1574         PyObject_GC_UnTrack(lz);
1575         Py_XDECREF(lz->ittuple);
1576         Py_XDECREF(lz->result);
1577         lz->ob_type->tp_free(lz);
1578 }
1579
1580 static int
1581 izip_traverse(izipobject *lz, visitproc visit, void *arg)
1582 {
1583         int err;
1584
1585         if (lz->ittuple) {
1586                 err = visit(lz->ittuple, arg);
1587                 if (err)
1588                         return err;
1589         }
1590         if (lz->result) {
1591                 err = visit(lz->result, arg);
1592                 if (err)
1593                         return err;
1594         }
1595         return 0;
1596 }
1597
1598 static PyObject *
1599 izip_next(izipobject *lz)
1600 {
1601         int i;
1602         long tuplesize = lz->tuplesize;
1603         PyObject *result = lz->result;
1604         PyObject *it;
1605         PyObject *item;
1606         PyObject *olditem;
1607
1608         if (tuplesize == 0)
1609                 return NULL;
1610         if (result->ob_refcnt == 1) {
1611                 Py_INCREF(result);
1612                 for (i=0 ; i < tuplesize ; i++) {
1613                         it = PyTuple_GET_ITEM(lz->ittuple, i);
1614                         assert(PyIter_Check(it));
1615                         item = (*it->ob_type->tp_iternext)(it);
1616                         if (item == NULL) {
1617                                 Py_DECREF(result);
1618                                 return NULL;
1619                         }
1620                         olditem = PyTuple_GET_ITEM(result, i);
1621                         PyTuple_SET_ITEM(result, i, item);
1622                         Py_DECREF(olditem);
1623                 }
1624         } else {
1625                 result = PyTuple_New(tuplesize);
1626                 if (result == NULL)
1627                         return NULL;
1628                 for (i=0 ; i < tuplesize ; i++) {
1629                         it = PyTuple_GET_ITEM(lz->ittuple, i);
1630                         assert(PyIter_Check(it));
1631                         item = (*it->ob_type->tp_iternext)(it);
1632                         if (item == NULL) {
1633                                 Py_DECREF(result);
1634                                 return NULL;
1635                         }
1636                         PyTuple_SET_ITEM(result, i, item);
1637                 }
1638         }
1639         return result;
1640 }
1641
1642 PyDoc_STRVAR(izip_doc,
1643 "izip(iter1 [,iter2 [...]]) --> izip object\n\
1644 \n\
1645 Return a izip object whose .next() method returns a tuple where\n\
1646 the i-th element comes from the i-th iterable argument.  The .next()\n\
1647 method continues until the shortest iterable in the argument sequence\n\
1648 is exhausted and then it raises StopIteration.  Works like the zip()\n\
1649 function but consumes less memory by returning an iterator instead of\n\
1650 a list.");
1651
1652 static PyTypeObject izip_type = {
1653         PyObject_HEAD_INIT(NULL)
1654         0,                              /* ob_size */
1655         "itertools.izip",               /* tp_name */
1656         sizeof(izipobject),             /* tp_basicsize */
1657         0,                              /* tp_itemsize */
1658         /* methods */
1659         (destructor)izip_dealloc,       /* tp_dealloc */
1660         0,                              /* tp_print */
1661         0,                              /* tp_getattr */
1662         0,                              /* tp_setattr */
1663         0,                              /* tp_compare */
1664         0,                              /* tp_repr */
1665         0,                              /* tp_as_number */
1666         0,                              /* tp_as_sequence */
1667         0,                              /* tp_as_mapping */
1668         0,                              /* tp_hash */
1669         0,                              /* tp_call */
1670         0,                              /* tp_str */
1671         PyObject_GenericGetAttr,        /* tp_getattro */
1672         0,                              /* tp_setattro */
1673         0,                              /* tp_as_buffer */
1674         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1675                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1676         izip_doc,                       /* tp_doc */
1677         (traverseproc)izip_traverse,    /* tp_traverse */
1678         0,                              /* tp_clear */
1679         0,                              /* tp_richcompare */
1680         0,                              /* tp_weaklistoffset */
1681         PyObject_SelfIter,              /* tp_iter */
1682         (iternextfunc)izip_next,        /* tp_iternext */
1683         0,                              /* tp_methods */
1684         0,                              /* tp_members */
1685         0,                              /* tp_getset */
1686         0,                              /* tp_base */
1687         0,                              /* tp_dict */
1688         0,                              /* tp_descr_get */
1689         0,                              /* tp_descr_set */
1690         0,                              /* tp_dictoffset */
1691         0,                              /* tp_init */
1692         0,                              /* tp_alloc */
1693         izip_new,                       /* tp_new */
1694         PyObject_GC_Del,                /* tp_free */
1695 };
1696
1697
1698 /* repeat object ************************************************************/
1699
1700 typedef struct {
1701         PyObject_HEAD
1702         PyObject *element;
1703         long cnt;
1704 } repeatobject;
1705
1706 static PyTypeObject repeat_type;
1707
1708 static PyObject *
1709 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1710 {
1711         repeatobject *ro;
1712         PyObject *element;
1713         long cnt = -1;
1714
1715         if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
1716                 return NULL;
1717
1718         if (PyTuple_Size(args) == 2 && cnt < 0)
1719                 cnt = 0;
1720
1721         ro = (repeatobject *)type->tp_alloc(type, 0);
1722         if (ro == NULL)
1723                 return NULL;
1724         Py_INCREF(element);
1725         ro->element = element;
1726         ro->cnt = cnt;
1727         return (PyObject *)ro;
1728 }
1729
1730 static void
1731 repeat_dealloc(repeatobject *ro)
1732 {
1733         PyObject_GC_UnTrack(ro);
1734         Py_XDECREF(ro->element);
1735         ro->ob_type->tp_free(ro);
1736 }
1737
1738 static int
1739 repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
1740 {
1741         if (ro->element)
1742                 return visit(ro->element, arg);
1743         return 0;
1744 }
1745
1746 static PyObject *
1747 repeat_next(repeatobject *ro)
1748 {
1749         if (ro->cnt == 0)
1750                 return NULL;
1751         if (ro->cnt > 0)
1752                 ro->cnt--;
1753         Py_INCREF(ro->element);
1754         return ro->element;
1755 }
1756
1757 PyDoc_STRVAR(repeat_doc,
1758 "repeat(element [,times]) -> create an iterator which returns the element\n\
1759 for the specified number of times.  If not specified, returns the element\n\
1760 endlessly.");
1761
1762 static PyTypeObject repeat_type = {
1763         PyObject_HEAD_INIT(NULL)
1764         0,                              /* ob_size */
1765         "itertools.repeat",             /* tp_name */
1766         sizeof(repeatobject),           /* tp_basicsize */
1767         0,                              /* tp_itemsize */
1768         /* methods */
1769         (destructor)repeat_dealloc,     /* tp_dealloc */
1770         0,                              /* tp_print */
1771         0,                              /* tp_getattr */
1772         0,                              /* tp_setattr */
1773         0,                              /* tp_compare */
1774         0,                              /* tp_repr */
1775         0,                              /* tp_as_number */
1776         0,                              /* tp_as_sequence */
1777         0,                              /* tp_as_mapping */
1778         0,                              /* tp_hash */
1779         0,                              /* tp_call */
1780         0,                              /* tp_str */
1781         PyObject_GenericGetAttr,        /* tp_getattro */
1782         0,                              /* tp_setattro */
1783         0,                              /* tp_as_buffer */
1784         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1785                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1786         repeat_doc,                     /* tp_doc */
1787         (traverseproc)repeat_traverse,  /* tp_traverse */
1788         0,                              /* tp_clear */
1789         0,                              /* tp_richcompare */
1790         0,                              /* tp_weaklistoffset */
1791         PyObject_SelfIter,              /* tp_iter */
1792         (iternextfunc)repeat_next,      /* tp_iternext */
1793         0,                              /* tp_methods */
1794         0,                              /* tp_members */
1795         0,                              /* tp_getset */
1796         0,                              /* tp_base */
1797         0,                              /* tp_dict */
1798         0,                              /* tp_descr_get */
1799         0,                              /* tp_descr_set */
1800         0,                              /* tp_dictoffset */
1801         0,                              /* tp_init */
1802         0,                              /* tp_alloc */
1803         repeat_new,                     /* tp_new */
1804         PyObject_GC_Del,                /* tp_free */
1805 };
1806
1807
1808 /* module level code ********************************************************/
1809
1810 PyDoc_STRVAR(module_doc,
1811 "Functional tools for creating and using iterators.\n\
1812 \n\
1813 Infinite iterators:\n\
1814 count([n]) --> n, n+1, n+2, ...\n\
1815 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
1816 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
1817 \n\
1818 Iterators terminating on the shortest input sequence:\n\
1819 izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
1820 ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
1821 ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
1822 islice(seq, [start,] stop [, step]) --> elements from\n\
1823        seq[start:stop:step]\n\
1824 imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
1825 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
1826 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
1827 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
1828 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
1829 ");
1830
1831
1832 PyMODINIT_FUNC
1833 inititertools(void)
1834 {
1835         int i;
1836         PyObject *m;
1837         char *name;
1838         PyTypeObject *typelist[] = {
1839                 &cycle_type,
1840                 &dropwhile_type,
1841                 &takewhile_type,
1842                 &islice_type,
1843                 &starmap_type,
1844                 &imap_type,
1845                 &chain_type,
1846                 &ifilter_type,
1847                 &ifilterfalse_type,
1848                 &count_type,
1849                 &izip_type,
1850                 &repeat_type,
1851                 NULL
1852         };
1853
1854         m = Py_InitModule3("itertools", NULL, module_doc);
1855
1856         for (i=0 ; typelist[i] != NULL ; i++) {
1857                 if (PyType_Ready(typelist[i]) < 0)
1858                         return;
1859                 name = strchr(typelist[i]->tp_name, '.');
1860                 assert (name != NULL);
1861                 Py_INCREF(typelist[i]);
1862                 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
1863         }
1864 }