Modules/itertoolsmodule.c

   1
   2 #include "Python.h"
   3 #include "structmember.h"
   4
   5 /* Itertools module written and maintained
   6    by Raymond D. Hettinger <python@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 const 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         gbo->ob_type->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         PyObject_HEAD_INIT(NULL)
 143         0,                              /* ob_size */
 144         "itertools.groupby",            /* tp_name */
 145         sizeof(groupbyobject),          /* tp_basicsize */
 146         0,                              /* tp_itemsize */
 147         /* methods */
 148         (destructor)groupby_dealloc,    /* tp_dealloc */
 149         0,                              /* tp_print */
 150         0,                              /* tp_getattr */
 151         0,                              /* tp_setattr */
 152         0,                              /* tp_compare */
 153         0,                              /* tp_repr */
 154         0,                              /* tp_as_number */
 155         0,                              /* tp_as_sequence */
 156         0,                              /* tp_as_mapping */
 157         0,                              /* tp_hash */
 158         0,                              /* tp_call */
 159         0,                              /* tp_str */
 160         PyObject_GenericGetAttr,        /* tp_getattro */
 161         0,                              /* tp_setattro */
 162         0,                              /* tp_as_buffer */
 163         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 164                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 165         groupby_doc,                    /* tp_doc */
 166         (traverseproc)groupby_traverse, /* tp_traverse */
 167         0,                              /* tp_clear */
 168         0,                              /* tp_richcompare */
 169         0,                              /* tp_weaklistoffset */
 170         PyObject_SelfIter,              /* tp_iter */
 171         (iternextfunc)groupby_next,     /* tp_iternext */
 172         0,                              /* tp_methods */
 173         0,                              /* tp_members */
 174         0,                              /* tp_getset */
 175         0,                              /* tp_base */
 176         0,                              /* tp_dict */
 177         0,                              /* tp_descr_get */
 178         0,                              /* tp_descr_set */
 179         0,                              /* tp_dictoffset */
 180         0,                              /* tp_init */
 181         0,                              /* tp_alloc */
 182         groupby_new,                    /* tp_new */
 183         PyObject_GC_Del,                /* tp_free */
 184 };
 185
 186
 187 /* _grouper object (internal) ************************************************/
 188
 189 typedef struct {
 190         PyObject_HEAD
 191         PyObject *parent;
 192         PyObject *tgtkey;
 193 } _grouperobject;
 194
 195 static PyTypeObject _grouper_type;
 196
 197 static PyObject *
 198 _grouper_create(groupbyobject *parent, PyObject *tgtkey)
 199 {
 200         _grouperobject *igo;
 201
 202         igo = PyObject_New(_grouperobject, &_grouper_type);
 203         if (igo == NULL)
 204                 return NULL;
 205         igo->parent = (PyObject *)parent;
 206         Py_INCREF(parent);
 207         igo->tgtkey = tgtkey;
 208         Py_INCREF(tgtkey);
 209
 210         return (PyObject *)igo;
 211 }
 212
 213 static void
 214 _grouper_dealloc(_grouperobject *igo)
 215 {
 216         Py_DECREF(igo->parent);
 217         Py_DECREF(igo->tgtkey);
 218         PyObject_Del(igo);
 219 }
 220
 221 static PyObject *
 222 _grouper_next(_grouperobject *igo)
 223 {
 224         groupbyobject *gbo = (groupbyobject *)igo->parent;
 225         PyObject *newvalue, *newkey, *r;
 226         int rcmp;
 227
 228         if (gbo->currvalue == NULL) {
 229                 newvalue = PyIter_Next(gbo->it);
 230                 if (newvalue == NULL)
 231                         return NULL;
 232
 233                 if (gbo->keyfunc == Py_None) {
 234                         newkey = newvalue;
 235                         Py_INCREF(newvalue);
 236                 } else {
 237                         newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
 238                                                               newvalue, NULL);
 239                         if (newkey == NULL) {
 240                                 Py_DECREF(newvalue);
 241                                 return NULL;
 242                         }
 243                 }
 244
 245                 assert(gbo->currkey == NULL);
 246                 gbo->currkey = newkey;
 247                 gbo->currvalue = newvalue;
 248         }
 249
 250         assert(gbo->currkey != NULL);
 251         rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
 252         if (rcmp <= 0)
 253                 /* got any error or current group is end */
 254                 return NULL;
 255
 256         r = gbo->currvalue;
 257         gbo->currvalue = NULL;
 258         Py_CLEAR(gbo->currkey);
 259
 260         return r;
 261 }
 262
 263 static PyTypeObject _grouper_type = {
 264         PyObject_HEAD_INIT(NULL)
 265         0,                              /* ob_size */
 266         "itertools._grouper",           /* tp_name */
 267         sizeof(_grouperobject),         /* tp_basicsize */
 268         0,                              /* tp_itemsize */
 269         /* methods */
 270         (destructor)_grouper_dealloc,   /* tp_dealloc */
 271         0,                              /* tp_print */
 272         0,                              /* tp_getattr */
 273         0,                              /* tp_setattr */
 274         0,                              /* tp_compare */
 275         0,                              /* tp_repr */
 276         0,                              /* tp_as_number */
 277         0,                              /* tp_as_sequence */
 278         0,                              /* tp_as_mapping */
 279         0,                              /* tp_hash */
 280         0,                              /* tp_call */
 281         0,                              /* tp_str */
 282         PyObject_GenericGetAttr,        /* tp_getattro */
 283         0,                              /* tp_setattro */
 284         0,                              /* tp_as_buffer */
 285         Py_TPFLAGS_DEFAULT,             /* tp_flags */
 286         0,                              /* tp_doc */
 287         0,                              /* tp_traverse */
 288         0,                              /* tp_clear */
 289         0,                              /* tp_richcompare */
 290         0,                              /* tp_weaklistoffset */
 291         PyObject_SelfIter,              /* tp_iter */
 292         (iternextfunc)_grouper_next,    /* tp_iternext */
 293         0,                              /* tp_methods */
 294         0,                              /* tp_members */
 295         0,                              /* tp_getset */
 296         0,                              /* tp_base */
 297         0,                              /* tp_dict */
 298         0,                              /* tp_descr_get */
 299         0,                              /* tp_descr_set */
 300         0,                              /* tp_dictoffset */
 301         0,                              /* tp_init */
 302         0,                              /* tp_alloc */
 303         0,                              /* tp_new */
 304         PyObject_Del,                   /* tp_free */
 305 };
 306
 307
 308
 309 /* tee object and with supporting function and objects ***************/
 310
 311 /* The teedataobject pre-allocates space for LINKCELLS number of objects.
 312    To help the object fit neatly inside cache lines (space for 16 to 32
 313    pointers), the value should be a multiple of 16 minus  space for
 314    the other structure members including PyHEAD overhead.  The larger the
 315    value, the less memory overhead per object and the less time spent
 316    allocating/deallocating new links.  The smaller the number, the less
 317    wasted space and the more rapid freeing of older data.
 318 */
 319 #define LINKCELLS 57
 320
 321 typedef struct {
 322         PyObject_HEAD
 323         PyObject *it;
 324         int numread;
 325         PyObject *nextlink;
 326         PyObject *(values[LINKCELLS]);
 327 } teedataobject;
 328
 329 typedef struct {
 330         PyObject_HEAD
 331         teedataobject *dataobj;
 332         int index;
 333         PyObject *weakreflist;
 334 } teeobject;
 335
 336 static PyTypeObject teedataobject_type;
 337
 338 static PyObject *
 339 teedataobject_new(PyObject *it)
 340 {
 341         teedataobject *tdo;
 342
 343         tdo = PyObject_New(teedataobject, &teedataobject_type);
 344         if (tdo == NULL)
 345                 return NULL;
 346
 347         tdo->numread = 0;
 348         tdo->nextlink = NULL;
 349         Py_INCREF(it);
 350         tdo->it = it;
 351         return (PyObject *)tdo;
 352 }
 353
 354 static PyObject *
 355 teedataobject_jumplink(teedataobject *tdo)
 356 {
 357         if (tdo->nextlink == NULL)
 358                 tdo->nextlink = teedataobject_new(tdo->it);
 359         Py_INCREF(tdo->nextlink);
 360         return tdo->nextlink;
 361 }
 362
 363 static PyObject *
 364 teedataobject_getitem(teedataobject *tdo, int i)
 365 {
 366         PyObject *value;
 367
 368         assert(i < LINKCELLS);
 369         if (i < tdo->numread)
 370                 value = tdo->values[i];
 371         else {
 372                 /* this is the lead iterator, so fetch more data */
 373                 assert(i == tdo->numread);
 374                 value = PyIter_Next(tdo->it);
 375                 if (value == NULL)
 376                         return NULL;
 377                 tdo->numread++;
 378                 tdo->values[i] = value;
 379         }
 380         Py_INCREF(value);
 381         return value;
 382 }
 383
 384 static void
 385 teedataobject_dealloc(teedataobject *tdo)
 386 {
 387         int i;
 388
 389         for (i=0 ; i<tdo->numread ; i++)
 390                 Py_DECREF(tdo->values[i]);
 391         Py_XDECREF(tdo->it);
 392         Py_XDECREF(tdo->nextlink);
 393         PyObject_Del(tdo);
 394 }
 395
 396 PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
 397
 398 static PyTypeObject teedataobject_type = {
 399         PyObject_HEAD_INIT(0)   /* Must fill in type value later */
 400         0,                                      /* ob_size */
 401         "itertools.tee_dataobject",             /* tp_name */
 402         sizeof(teedataobject),                  /* tp_basicsize */
 403         0,                                      /* tp_itemsize */
 404         /* methods */
 405         (destructor)teedataobject_dealloc,      /* tp_dealloc */
 406         0,                                      /* tp_print */
 407         0,                                      /* tp_getattr */
 408         0,                                      /* tp_setattr */
 409         0,                                      /* tp_compare */
 410         0,                                      /* tp_repr */
 411         0,                                      /* tp_as_number */
 412         0,                                      /* tp_as_sequence */
 413         0,                                      /* tp_as_mapping */
 414         0,                                      /* tp_hash */
 415         0,                                      /* tp_call */
 416         0,                                      /* tp_str */
 417         PyObject_GenericGetAttr,                /* tp_getattro */
 418         0,                                      /* tp_setattro */
 419         0,                                      /* tp_as_buffer */
 420         Py_TPFLAGS_DEFAULT,                     /* tp_flags */
 421         teedataobject_doc,                      /* tp_doc */
 422         0,                                      /* tp_traverse */
 423 };
 424
 425
 426 static PyTypeObject tee_type;
 427
 428 static PyObject *
 429 tee_next(teeobject *to)
 430 {
 431         PyObject *value, *link;
 432
 433         if (to->index >= LINKCELLS) {
 434                 link = teedataobject_jumplink(to->dataobj);
 435                 Py_XDECREF(to->dataobj);
 436                 to->dataobj = (teedataobject *)link;
 437                 to->index = 0;
 438         }
 439         value = teedataobject_getitem(to->dataobj, to->index);
 440         if (value == NULL)
 441                 return NULL;
 442         to->index++;
 443         return value;
 444 }
 445
 446 static PyObject *
 447 tee_copy(teeobject *to)
 448 {
 449         teeobject *newto;
 450
 451         newto = PyObject_New(teeobject, &tee_type);
 452         if (newto == NULL)
 453                 return NULL;
 454         Py_INCREF(to->dataobj);
 455         newto->dataobj = to->dataobj;
 456         newto->index = to->index;
 457         newto->weakreflist = NULL;
 458         return (PyObject *)newto;
 459 }
 460
 461 PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
 462
 463 static PyObject *
 464 tee_fromiterable(PyObject *iterable)
 465 {
 466         teeobject *to;
 467         PyObject *it = NULL;
 468
 469         it = PyObject_GetIter(iterable);
 470         if (it == NULL)
 471                 return NULL;
 472         if (PyObject_TypeCheck(it, &tee_type)) {
 473                 to = (teeobject *)tee_copy((teeobject *)it);
 474                 goto done;
 475         }
 476
 477         to = PyObject_New(teeobject, &tee_type);
 478         if (to == NULL)
 479                 goto done;
 480         to->dataobj = (teedataobject *)teedataobject_new(it);
 481         to->index = 0;
 482         to->weakreflist = NULL;
 483 done:
 484         Py_XDECREF(it);
 485         return (PyObject *)to;
 486 }
 487
 488 static PyObject *
 489 tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 490 {
 491         PyObject *iterable;
 492
 493         if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
 494                 return NULL;
 495         return tee_fromiterable(iterable);
 496 }
 497
 498 static void
 499 tee_dealloc(teeobject *to)
 500 {
 501         if (to->weakreflist != NULL)
 502                 PyObject_ClearWeakRefs((PyObject *) to);
 503         Py_XDECREF(to->dataobj);
 504         PyObject_Del(to);
 505 }
 506
 507 PyDoc_STRVAR(teeobject_doc,
 508 "Iterator wrapped to make it copyable");
 509
 510 static PyMethodDef tee_methods[] = {
 511         {"__copy__",    (PyCFunction)tee_copy,  METH_NOARGS, teecopy_doc},
 512         {NULL,          NULL}           /* sentinel */
 513 };
 514
 515 static PyTypeObject tee_type = {
 516         PyObject_HEAD_INIT(NULL)
 517         0,                              /* ob_size */
 518         "itertools.tee",                /* tp_name */
 519         sizeof(teeobject),              /* tp_basicsize */
 520         0,                              /* tp_itemsize */
 521         /* methods */
 522         (destructor)tee_dealloc,        /* tp_dealloc */
 523         0,                              /* tp_print */
 524         0,                              /* tp_getattr */
 525         0,                              /* tp_setattr */
 526         0,                              /* tp_compare */
 527         0,                              /* tp_repr */
 528         0,                              /* tp_as_number */
 529         0,                              /* tp_as_sequence */
 530         0,                              /* tp_as_mapping */
 531         0,                              /* tp_hash */
 532         0,                              /* tp_call */
 533         0,                              /* tp_str */
 534         0,                              /* tp_getattro */
 535         0,                              /* tp_setattro */
 536         0,                              /* tp_as_buffer */
 537         Py_TPFLAGS_DEFAULT,             /* tp_flags */
 538         teeobject_doc,                  /* tp_doc */
 539         0,                              /* tp_traverse */
 540         0,                              /* tp_clear */
 541         0,                              /* tp_richcompare */
 542         offsetof(teeobject, weakreflist),       /* tp_weaklistoffset */
 543         PyObject_SelfIter,              /* tp_iter */
 544         (iternextfunc)tee_next,         /* tp_iternext */
 545         tee_methods,                    /* tp_methods */
 546         0,                              /* tp_members */
 547         0,                              /* tp_getset */
 548         0,                              /* tp_base */
 549         0,                              /* tp_dict */
 550         0,                              /* tp_descr_get */
 551         0,                              /* tp_descr_set */
 552         0,                              /* tp_dictoffset */
 553         0,                              /* tp_init */
 554         0,                              /* tp_alloc */
 555         tee_new,                        /* tp_new */
 556         PyObject_Del,                   /* tp_free */
 557 };
 558
 559 static PyObject *
 560 tee(PyObject *self, PyObject *args)
 561 {
 562         int i, n=2;
 563         PyObject *it, *iterable, *copyable, *result;
 564
 565         if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
 566                 return NULL;
 567         result = PyTuple_New(n);
 568         if (result == NULL)
 569                 return NULL;
 570         if (n == 0)
 571                 return result;
 572         it = PyObject_GetIter(iterable);
 573         if (it == NULL) {
 574                 Py_DECREF(result);
 575                 return NULL;
 576         }
 577         if (!PyObject_HasAttrString(it, "__copy__")) {
 578                 copyable = tee_fromiterable(it);
 579                 Py_DECREF(it);
 580                 if (copyable == NULL) {
 581                         Py_DECREF(result);
 582                         return NULL;
 583                 }
 584         } else
 585                 copyable = it;
 586         PyTuple_SET_ITEM(result, 0, copyable);
 587         for (i=1 ; i<n ; i++) {
 588                 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
 589                 if (copyable == NULL) {
 590                         Py_DECREF(result);
 591                         return NULL;
 592                 }
 593                 PyTuple_SET_ITEM(result, i, copyable);
 594         }
 595         return result;
 596 }
 597
 598 PyDoc_STRVAR(tee_doc,
 599 "tee(iterable, n=2) --> tuple of n independent iterators.");
 600
 601
 602 /* cycle object **********************************************************/
 603
 604 typedef struct {
 605         PyObject_HEAD
 606         PyObject *it;
 607         PyObject *saved;
 608         int firstpass;
 609 } cycleobject;
 610
 611 static PyTypeObject cycle_type;
 612
 613 static PyObject *
 614 cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 615 {
 616         PyObject *it;
 617         PyObject *iterable;
 618         PyObject *saved;
 619         cycleobject *lz;
 620
 621         if (!_PyArg_NoKeywords("cycle()", kwds))
 622                 return NULL;
 623
 624         if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
 625                 return NULL;
 626
 627         /* Get iterator. */
 628         it = PyObject_GetIter(iterable);
 629         if (it == NULL)
 630                 return NULL;
 631
 632         saved = PyList_New(0);
 633         if (saved == NULL) {
 634                 Py_DECREF(it);
 635                 return NULL;
 636         }
 637
 638         /* create cycleobject structure */
 639         lz = (cycleobject *)type->tp_alloc(type, 0);
 640         if (lz == NULL) {
 641                 Py_DECREF(it);
 642                 Py_DECREF(saved);
 643                 return NULL;
 644         }
 645         lz->it = it;
 646         lz->saved = saved;
 647         lz->firstpass = 0;
 648
 649         return (PyObject *)lz;
 650 }
 651
 652 static void
 653 cycle_dealloc(cycleobject *lz)
 654 {
 655         PyObject_GC_UnTrack(lz);
 656         Py_XDECREF(lz->saved);
 657         Py_XDECREF(lz->it);
 658         lz->ob_type->tp_free(lz);
 659 }
 660
 661 static int
 662 cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
 663 {
 664         Py_VISIT(lz->it);
 665         Py_VISIT(lz->saved);
 666         return 0;
 667 }
 668
 669 static PyObject *
 670 cycle_next(cycleobject *lz)
 671 {
 672         PyObject *item;
 673         PyObject *it;
 674         PyObject *tmp;
 675
 676         while (1) {
 677                 item = PyIter_Next(lz->it);
 678                 if (item != NULL) {
 679                         if (!lz->firstpass)
 680                                 PyList_Append(lz->saved, item);
 681                         return item;
 682                 }
 683                 if (PyErr_Occurred()) {
 684                         if (PyErr_ExceptionMatches(PyExc_StopIteration))
 685                                 PyErr_Clear();
 686                         else
 687                                 return NULL;
 688                 }
 689                 if (PyList_Size(lz->saved) == 0)
 690                         return NULL;
 691                 it = PyObject_GetIter(lz->saved);
 692                 if (it == NULL)
 693                         return NULL;
 694                 tmp = lz->it;
 695                 lz->it = it;
 696                 lz->firstpass = 1;
 697                 Py_DECREF(tmp);
 698         }
 699 }
 700
 701 PyDoc_STRVAR(cycle_doc,
 702 "cycle(iterable) --> cycle object\n\
 703 \n\
 704 Return elements from the iterable until it is exhausted.\n\
 705 Then repeat the sequence indefinitely.");
 706
 707 static PyTypeObject cycle_type = {
 708         PyObject_HEAD_INIT(NULL)
 709         0,                              /* ob_size */
 710         "itertools.cycle",              /* tp_name */
 711         sizeof(cycleobject),            /* tp_basicsize */
 712         0,                              /* tp_itemsize */
 713         /* methods */
 714         (destructor)cycle_dealloc,      /* tp_dealloc */
 715         0,                              /* tp_print */
 716         0,                              /* tp_getattr */
 717         0,                              /* tp_setattr */
 718         0,                              /* tp_compare */
 719         0,                              /* tp_repr */
 720         0,                              /* tp_as_number */
 721         0,                              /* tp_as_sequence */
 722         0,                              /* tp_as_mapping */
 723         0,                              /* tp_hash */
 724         0,                              /* tp_call */
 725         0,                              /* tp_str */
 726         PyObject_GenericGetAttr,        /* tp_getattro */
 727         0,                              /* tp_setattro */
 728         0,                              /* tp_as_buffer */
 729         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 730                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 731         cycle_doc,                      /* tp_doc */
 732         (traverseproc)cycle_traverse,   /* tp_traverse */
 733         0,                              /* tp_clear */
 734         0,                              /* tp_richcompare */
 735         0,                              /* tp_weaklistoffset */
 736         PyObject_SelfIter,              /* tp_iter */
 737         (iternextfunc)cycle_next,       /* tp_iternext */
 738         0,                              /* tp_methods */
 739         0,                              /* tp_members */
 740         0,                              /* tp_getset */
 741         0,                              /* tp_base */
 742         0,                              /* tp_dict */
 743         0,                              /* tp_descr_get */
 744         0,                              /* tp_descr_set */
 745         0,                              /* tp_dictoffset */
 746         0,                              /* tp_init */
 747         0,                              /* tp_alloc */
 748         cycle_new,                      /* tp_new */
 749         PyObject_GC_Del,                /* tp_free */
 750 };
 751
 752
 753 /* dropwhile object **********************************************************/
 754
 755 typedef struct {
 756         PyObject_HEAD
 757         PyObject *func;
 758         PyObject *it;
 759         long     start;
 760 } dropwhileobject;
 761
 762 static PyTypeObject dropwhile_type;
 763
 764 static PyObject *
 765 dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 766 {
 767         PyObject *func, *seq;
 768         PyObject *it;
 769         dropwhileobject *lz;
 770
 771         if (!_PyArg_NoKeywords("dropwhile()", kwds))
 772                 return NULL;
 773
 774         if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
 775                 return NULL;
 776
 777         /* Get iterator. */
 778         it = PyObject_GetIter(seq);
 779         if (it == NULL)
 780                 return NULL;
 781
 782         /* create dropwhileobject structure */
 783         lz = (dropwhileobject *)type->tp_alloc(type, 0);
 784         if (lz == NULL) {
 785                 Py_DECREF(it);
 786                 return NULL;
 787         }
 788         Py_INCREF(func);
 789         lz->func = func;
 790         lz->it = it;
 791         lz->start = 0;
 792
 793         return (PyObject *)lz;
 794 }
 795
 796 static void
 797 dropwhile_dealloc(dropwhileobject *lz)
 798 {
 799         PyObject_GC_UnTrack(lz);
 800         Py_XDECREF(lz->func);
 801         Py_XDECREF(lz->it);
 802         lz->ob_type->tp_free(lz);
 803 }
 804
 805 static int
 806 dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
 807 {
 808         Py_VISIT(lz->it);
 809         Py_VISIT(lz->func);
 810         return 0;
 811 }
 812
 813 static PyObject *
 814 dropwhile_next(dropwhileobject *lz)
 815 {
 816         PyObject *item, *good;
 817         PyObject *it = lz->it;
 818         long ok;
 819         PyObject *(*iternext)(PyObject *);
 820
 821         assert(PyIter_Check(it));
 822         iternext = *it->ob_type->tp_iternext;
 823         for (;;) {
 824                 item = iternext(it);
 825                 if (item == NULL)
 826                         return NULL;
 827                 if (lz->start == 1)
 828                         return item;
 829
 830                 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
 831                 if (good == NULL) {
 832                         Py_DECREF(item);
 833                         return NULL;
 834                 }
 835                 ok = PyObject_IsTrue(good);
 836                 Py_DECREF(good);
 837                 if (!ok) {
 838                         lz->start = 1;
 839                         return item;
 840                 }
 841                 Py_DECREF(item);
 842         }
 843 }
 844
 845 PyDoc_STRVAR(dropwhile_doc,
 846 "dropwhile(predicate, iterable) --> dropwhile object\n\
 847 \n\
 848 Drop items from the iterable while predicate(item) is true.\n\
 849 Afterwards, return every element until the iterable is exhausted.");
 850
 851 static PyTypeObject dropwhile_type = {
 852         PyObject_HEAD_INIT(NULL)
 853         0,                              /* ob_size */
 854         "itertools.dropwhile",          /* tp_name */
 855         sizeof(dropwhileobject),        /* tp_basicsize */
 856         0,                              /* tp_itemsize */
 857         /* methods */
 858         (destructor)dropwhile_dealloc,  /* tp_dealloc */
 859         0,                              /* tp_print */
 860         0,                              /* tp_getattr */
 861         0,                              /* tp_setattr */
 862         0,                              /* tp_compare */
 863         0,                              /* tp_repr */
 864         0,                              /* tp_as_number */
 865         0,                              /* tp_as_sequence */
 866         0,                              /* tp_as_mapping */
 867         0,                              /* tp_hash */
 868         0,                              /* tp_call */
 869         0,                              /* tp_str */
 870         PyObject_GenericGetAttr,        /* tp_getattro */
 871         0,                              /* tp_setattro */
 872         0,                              /* tp_as_buffer */
 873         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 874                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 875         dropwhile_doc,                  /* tp_doc */
 876         (traverseproc)dropwhile_traverse,    /* tp_traverse */
 877         0,                              /* tp_clear */
 878         0,                              /* tp_richcompare */
 879         0,                              /* tp_weaklistoffset */
 880         PyObject_SelfIter,              /* tp_iter */
 881         (iternextfunc)dropwhile_next,   /* tp_iternext */
 882         0,                              /* tp_methods */
 883         0,                              /* tp_members */
 884         0,                              /* tp_getset */
 885         0,                              /* tp_base */
 886         0,                              /* tp_dict */
 887         0,                              /* tp_descr_get */
 888         0,                              /* tp_descr_set */
 889         0,                              /* tp_dictoffset */
 890         0,                              /* tp_init */
 891         0,                              /* tp_alloc */
 892         dropwhile_new,                  /* tp_new */
 893         PyObject_GC_Del,                /* tp_free */
 894 };
 895
 896
 897 /* takewhile object **********************************************************/
 898
 899 typedef struct {
 900         PyObject_HEAD
 901         PyObject *func;
 902         PyObject *it;
 903         long     stop;
 904 } takewhileobject;
 905
 906 static PyTypeObject takewhile_type;
 907
 908 static PyObject *
 909 takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 910 {
 911         PyObject *func, *seq;
 912         PyObject *it;
 913         takewhileobject *lz;
 914
 915         if (!_PyArg_NoKeywords("takewhile()", kwds))
 916                 return NULL;
 917
 918         if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
 919                 return NULL;
 920
 921         /* Get iterator. */
 922         it = PyObject_GetIter(seq);
 923         if (it == NULL)
 924                 return NULL;
 925
 926         /* create takewhileobject structure */
 927         lz = (takewhileobject *)type->tp_alloc(type, 0);
 928         if (lz == NULL) {
 929                 Py_DECREF(it);
 930                 return NULL;
 931         }
 932         Py_INCREF(func);
 933         lz->func = func;
 934         lz->it = it;
 935         lz->stop = 0;
 936
 937         return (PyObject *)lz;
 938 }
 939
 940 static void
 941 takewhile_dealloc(takewhileobject *lz)
 942 {
 943         PyObject_GC_UnTrack(lz);
 944         Py_XDECREF(lz->func);
 945         Py_XDECREF(lz->it);
 946         lz->ob_type->tp_free(lz);
 947 }
 948
 949 static int
 950 takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
 951 {
 952         Py_VISIT(lz->it);
 953         Py_VISIT(lz->func);
 954         return 0;
 955 }
 956
 957 static PyObject *
 958 takewhile_next(takewhileobject *lz)
 959 {
 960         PyObject *item, *good;
 961         PyObject *it = lz->it;
 962         long ok;
 963
 964         if (lz->stop == 1)
 965                 return NULL;
 966
 967         assert(PyIter_Check(it));
 968         item = (*it->ob_type->tp_iternext)(it);
 969         if (item == NULL)
 970                 return NULL;
 971
 972         good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
 973         if (good == NULL) {
 974                 Py_DECREF(item);
 975                 return NULL;
 976         }
 977         ok = PyObject_IsTrue(good);
 978         Py_DECREF(good);
 979         if (ok)
 980                 return item;
 981         Py_DECREF(item);
 982         lz->stop = 1;
 983         return NULL;
 984 }
 985
 986 PyDoc_STRVAR(takewhile_doc,
 987 "takewhile(predicate, iterable) --> takewhile object\n\
 988 \n\
 989 Return successive entries from an iterable as long as the \n\
 990 predicate evaluates to true for each entry.");
 991
 992 static PyTypeObject takewhile_type = {
 993         PyObject_HEAD_INIT(NULL)
 994         0,                              /* ob_size */
 995         "itertools.takewhile",          /* tp_name */
 996         sizeof(takewhileobject),        /* tp_basicsize */
 997         0,                              /* tp_itemsize */
 998         /* methods */
 999         (destructor)takewhile_dealloc,  /* tp_dealloc */
1000         0,                              /* tp_print */
1001         0,                              /* tp_getattr */
1002         0,                              /* tp_setattr */
1003         0,                              /* tp_compare */
1004         0,                              /* tp_repr */
1005         0,                              /* tp_as_number */
1006         0,                              /* tp_as_sequence */
1007         0,                              /* tp_as_mapping */
1008         0,                              /* tp_hash */
1009         0,                              /* tp_call */
1010         0,                              /* tp_str */
1011         PyObject_GenericGetAttr,        /* tp_getattro */
1012         0,                              /* tp_setattro */
1013         0,                              /* tp_as_buffer */
1014         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1015                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1016         takewhile_doc,                  /* tp_doc */
1017         (traverseproc)takewhile_traverse,    /* tp_traverse */
1018         0,                              /* tp_clear */
1019         0,                              /* tp_richcompare */
1020         0,                              /* tp_weaklistoffset */
1021         PyObject_SelfIter,              /* tp_iter */
1022         (iternextfunc)takewhile_next,   /* tp_iternext */
1023         0,                              /* tp_methods */
1024         0,                              /* tp_members */
1025         0,                              /* tp_getset */
1026         0,                              /* tp_base */
1027         0,                              /* tp_dict */
1028         0,                              /* tp_descr_get */
1029         0,                              /* tp_descr_set */
1030         0,                              /* tp_dictoffset */
1031         0,                              /* tp_init */
1032         0,                              /* tp_alloc */
1033         takewhile_new,                  /* tp_new */
1034         PyObject_GC_Del,                /* tp_free */
1035 };
1036
1037
1038 /* islice object ************************************************************/
1039
1040 typedef struct {
1041         PyObject_HEAD
1042         PyObject *it;
1043         long    next;
1044         long    stop;
1045         long    step;
1046         long    cnt;
1047 } isliceobject;
1048
1049 static PyTypeObject islice_type;
1050
1051 static PyObject *
1052 islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1053 {
1054         PyObject *seq;
1055         long start=0, stop=-1, step=1;
1056         PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1057         int numargs;
1058         isliceobject *lz;
1059
1060         if (!_PyArg_NoKeywords("islice()", kwds))
1061                 return NULL;
1062
1063         if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1064                 return NULL;
1065
1066         numargs = PyTuple_Size(args);
1067         if (numargs == 2) {
1068                 if (a1 != Py_None) {
1069                         stop = PyInt_AsLong(a1);
1070                         if (stop == -1) {
1071                                 if (PyErr_Occurred())
1072                                         PyErr_Clear();
1073                                 PyErr_SetString(PyExc_ValueError,
1074                                    "Stop argument for islice() must be a non-negative integer or None.");
1075                                 return NULL;
1076                         }
1077                 }
1078         } else {
1079                 if (a1 != Py_None)
1080                         start = PyInt_AsLong(a1);
1081                 if (start == -1 && PyErr_Occurred())
1082                         PyErr_Clear();
1083                 if (a2 != Py_None) {
1084                         stop = PyInt_AsLong(a2);
1085                         if (stop == -1) {
1086                                 if (PyErr_Occurred())
1087                                         PyErr_Clear();
1088                                 PyErr_SetString(PyExc_ValueError,
1089                                    "Stop argument for islice() must be a non-negative integer or None.");
1090                                 return NULL;
1091                         }
1092                 }
1093         }
1094         if (start<0 || stop<-1) {
1095                 PyErr_SetString(PyExc_ValueError,
1096                    "Indices for islice() must be non-negative integers or None.");
1097                 return NULL;
1098         }
1099
1100         if (a3 != NULL) {
1101                 if (a3 != Py_None)
1102                         step = PyInt_AsLong(a3);
1103                 if (step == -1 && PyErr_Occurred())
1104                         PyErr_Clear();
1105         }
1106         if (step<1) {
1107                 PyErr_SetString(PyExc_ValueError,
1108                    "Step for islice() must be a positive integer or None.");
1109                 return NULL;
1110         }
1111
1112         /* Get iterator. */
1113         it = PyObject_GetIter(seq);
1114         if (it == NULL)
1115                 return NULL;
1116
1117         /* create isliceobject structure */
1118         lz = (isliceobject *)type->tp_alloc(type, 0);
1119         if (lz == NULL) {
1120                 Py_DECREF(it);
1121                 return NULL;
1122         }
1123         lz->it = it;
1124         lz->next = start;
1125         lz->stop = stop;
1126         lz->step = step;
1127         lz->cnt = 0L;
1128
1129         return (PyObject *)lz;
1130 }
1131
1132 static void
1133 islice_dealloc(isliceobject *lz)
1134 {
1135         PyObject_GC_UnTrack(lz);
1136         Py_XDECREF(lz->it);
1137         lz->ob_type->tp_free(lz);
1138 }
1139
1140 static int
1141 islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1142 {
1143         Py_VISIT(lz->it);
1144         return 0;
1145 }
1146
1147 static PyObject *
1148 islice_next(isliceobject *lz)
1149 {
1150         PyObject *item;
1151         PyObject *it = lz->it;
1152         long oldnext;
1153         PyObject *(*iternext)(PyObject *);
1154
1155         assert(PyIter_Check(it));
1156         iternext = *it->ob_type->tp_iternext;
1157         while (lz->cnt < lz->next) {
1158                 item = iternext(it);
1159                 if (item == NULL)
1160                         return NULL;
1161                 Py_DECREF(item);
1162                 lz->cnt++;
1163         }
1164         if (lz->stop != -1 && lz->cnt >= lz->stop)
1165                 return NULL;
1166         assert(PyIter_Check(it));
1167         item = iternext(it);
1168         if (item == NULL)
1169                 return NULL;
1170         lz->cnt++;
1171         oldnext = lz->next;
1172         lz->next += lz->step;
1173         if (lz->next < oldnext) /* Check for overflow */
1174                 lz->next = lz->stop;
1175         return item;
1176 }
1177
1178 PyDoc_STRVAR(islice_doc,
1179 "islice(iterable, [start,] stop [, step]) --> islice object\n\
1180 \n\
1181 Return an iterator whose next() method returns selected values from an\n\
1182 iterable.  If start is specified, will skip all preceding elements;\n\
1183 otherwise, start defaults to zero.  Step defaults to one.  If\n\
1184 specified as another value, step determines how many values are \n\
1185 skipped between successive calls.  Works like a slice() on a list\n\
1186 but returns an iterator.");
1187
1188 static PyTypeObject islice_type = {
1189         PyObject_HEAD_INIT(NULL)
1190         0,                              /* ob_size */
1191         "itertools.islice",             /* tp_name */
1192         sizeof(isliceobject),           /* tp_basicsize */
1193         0,                              /* tp_itemsize */
1194         /* methods */
1195         (destructor)islice_dealloc,     /* tp_dealloc */
1196         0,                              /* tp_print */
1197         0,                              /* tp_getattr */
1198         0,                              /* tp_setattr */
1199         0,                              /* tp_compare */
1200         0,                              /* tp_repr */
1201         0,                              /* tp_as_number */
1202         0,                              /* tp_as_sequence */
1203         0,                              /* tp_as_mapping */
1204         0,                              /* tp_hash */
1205         0,                              /* tp_call */
1206         0,                              /* tp_str */
1207         PyObject_GenericGetAttr,        /* tp_getattro */
1208         0,                              /* tp_setattro */
1209         0,                              /* tp_as_buffer */
1210         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1211                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1212         islice_doc,                     /* tp_doc */
1213         (traverseproc)islice_traverse,  /* tp_traverse */
1214         0,                              /* tp_clear */
1215         0,                              /* tp_richcompare */
1216         0,                              /* tp_weaklistoffset */
1217         PyObject_SelfIter,              /* tp_iter */
1218         (iternextfunc)islice_next,      /* tp_iternext */
1219         0,                              /* tp_methods */
1220         0,                              /* tp_members */
1221         0,                              /* tp_getset */
1222         0,                              /* tp_base */
1223         0,                              /* tp_dict */
1224         0,                              /* tp_descr_get */
1225         0,                              /* tp_descr_set */
1226         0,                              /* tp_dictoffset */
1227         0,                              /* tp_init */
1228         0,                              /* tp_alloc */
1229         islice_new,                     /* tp_new */
1230         PyObject_GC_Del,                /* tp_free */
1231 };
1232
1233
1234 /* starmap object ************************************************************/
1235
1236 typedef struct {
1237         PyObject_HEAD
1238         PyObject *func;
1239         PyObject *it;
1240 } starmapobject;
1241
1242 static PyTypeObject starmap_type;
1243
1244 static PyObject *
1245 starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1246 {
1247         PyObject *func, *seq;
1248         PyObject *it;
1249         starmapobject *lz;
1250
1251         if (!_PyArg_NoKeywords("starmap()", kwds))
1252                 return NULL;
1253
1254         if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1255                 return NULL;
1256
1257         /* Get iterator. */
1258         it = PyObject_GetIter(seq);
1259         if (it == NULL)
1260                 return NULL;
1261
1262         /* create starmapobject structure */
1263         lz = (starmapobject *)type->tp_alloc(type, 0);
1264         if (lz == NULL) {
1265                 Py_DECREF(it);
1266                 return NULL;
1267         }
1268         Py_INCREF(func);
1269         lz->func = func;
1270         lz->it = it;
1271
1272         return (PyObject *)lz;
1273 }
1274
1275 static void
1276 starmap_dealloc(starmapobject *lz)
1277 {
1278         PyObject_GC_UnTrack(lz);
1279         Py_XDECREF(lz->func);
1280         Py_XDECREF(lz->it);
1281         lz->ob_type->tp_free(lz);
1282 }
1283
1284 static int
1285 starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1286 {
1287         Py_VISIT(lz->it);
1288         Py_VISIT(lz->func);
1289         return 0;
1290 }
1291
1292 static PyObject *
1293 starmap_next(starmapobject *lz)
1294 {
1295         PyObject *args;
1296         PyObject *result;
1297         PyObject *it = lz->it;
1298
1299         assert(PyIter_Check(it));
1300         args = (*it->ob_type->tp_iternext)(it);
1301         if (args == NULL)
1302                 return NULL;
1303         if (!PyTuple_CheckExact(args)) {
1304                 Py_DECREF(args);
1305                 PyErr_SetString(PyExc_TypeError,
1306                                 "iterator must return a tuple");
1307                 return NULL;
1308         }
1309         result = PyObject_Call(lz->func, args, NULL);
1310         Py_DECREF(args);
1311         return result;
1312 }
1313
1314 PyDoc_STRVAR(starmap_doc,
1315 "starmap(function, sequence) --> starmap object\n\
1316 \n\
1317 Return an iterator whose values are returned from the function evaluated\n\
1318 with a argument tuple taken from the given sequence.");
1319
1320 static PyTypeObject starmap_type = {
1321         PyObject_HEAD_INIT(NULL)
1322         0,                              /* ob_size */
1323         "itertools.starmap",            /* tp_name */
1324         sizeof(starmapobject),          /* tp_basicsize */
1325         0,                              /* tp_itemsize */
1326         /* methods */
1327         (destructor)starmap_dealloc,    /* tp_dealloc */
1328         0,                              /* tp_print */
1329         0,                              /* tp_getattr */
1330         0,                              /* tp_setattr */
1331         0,                              /* tp_compare */
1332         0,                              /* tp_repr */
1333         0,                              /* tp_as_number */
1334         0,                              /* tp_as_sequence */
1335         0,                              /* tp_as_mapping */
1336         0,                              /* tp_hash */
1337         0,                              /* tp_call */
1338         0,                              /* tp_str */
1339         PyObject_GenericGetAttr,        /* tp_getattro */
1340         0,                              /* tp_setattro */
1341         0,                              /* tp_as_buffer */
1342         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1343                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1344         starmap_doc,                    /* tp_doc */
1345         (traverseproc)starmap_traverse, /* tp_traverse */
1346         0,                              /* tp_clear */
1347         0,                              /* tp_richcompare */
1348         0,                              /* tp_weaklistoffset */
1349         PyObject_SelfIter,              /* tp_iter */
1350         (iternextfunc)starmap_next,     /* tp_iternext */
1351         0,                              /* tp_methods */
1352         0,                              /* tp_members */
1353         0,                              /* tp_getset */
1354         0,                              /* tp_base */
1355         0,                              /* tp_dict */
1356         0,                              /* tp_descr_get */
1357         0,                              /* tp_descr_set */
1358         0,                              /* tp_dictoffset */
1359         0,                              /* tp_init */
1360         0,                              /* tp_alloc */
1361         starmap_new,                    /* tp_new */
1362         PyObject_GC_Del,                /* tp_free */
1363 };
1364
1365
1366 /* imap object ************************************************************/
1367
1368 typedef struct {
1369         PyObject_HEAD
1370         PyObject *iters;
1371         PyObject *func;
1372 } imapobject;
1373
1374 static PyTypeObject imap_type;
1375
1376 static PyObject *
1377 imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1378 {
1379         PyObject *it, *iters, *func;
1380         imapobject *lz;
1381         int numargs, i;
1382
1383         if (!_PyArg_NoKeywords("imap()", kwds))
1384                 return NULL;
1385
1386         numargs = PyTuple_Size(args);
1387         if (numargs < 2) {
1388                 PyErr_SetString(PyExc_TypeError,
1389                    "imap() must have at least two arguments.");
1390                 return NULL;
1391         }
1392
1393         iters = PyTuple_New(numargs-1);
1394         if (iters == NULL)
1395                 return NULL;
1396
1397         for (i=1 ; i<numargs ; i++) {
1398                 /* Get iterator. */
1399                 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1400                 if (it == NULL) {
1401                         Py_DECREF(iters);
1402                         return NULL;
1403                 }
1404                 PyTuple_SET_ITEM(iters, i-1, it);
1405         }
1406
1407         /* create imapobject structure */
1408         lz = (imapobject *)type->tp_alloc(type, 0);
1409         if (lz == NULL) {
1410                 Py_DECREF(iters);
1411                 return NULL;
1412         }
1413         lz->iters = iters;
1414         func = PyTuple_GET_ITEM(args, 0);
1415         Py_INCREF(func);
1416         lz->func = func;
1417
1418         return (PyObject *)lz;
1419 }
1420
1421 static void
1422 imap_dealloc(imapobject *lz)
1423 {
1424         PyObject_GC_UnTrack(lz);
1425         Py_XDECREF(lz->iters);
1426         Py_XDECREF(lz->func);
1427         lz->ob_type->tp_free(lz);
1428 }
1429
1430 static int
1431 imap_traverse(imapobject *lz, visitproc visit, void *arg)
1432 {
1433         Py_VISIT(lz->iters);
1434         Py_VISIT(lz->func);
1435         return 0;
1436 }
1437
1438 /*
1439 imap() is an iterator version of __builtins__.map() except that it does
1440 not have the None fill-in feature.  That was intentionally left out for
1441 the following reasons:
1442
1443   1) Itertools are designed to be easily combined and chained together.
1444      Having all tools stop with the shortest input is a unifying principle
1445      that makes it easier to combine finite iterators (supplying data) with
1446      infinite iterators like count() and repeat() (for supplying sequential
1447      or constant arguments to a function).
1448
1449   2) In typical use cases for combining itertools, having one finite data
1450      supplier run out before another is likely to be an error condition which
1451      should not pass silently by automatically supplying None.
1452
1453   3) The use cases for automatic None fill-in are rare -- not many functions
1454      do something useful when a parameter suddenly switches type and becomes
1455      None.
1456
1457   4) If a need does arise, it can be met by __builtins__.map() or by
1458      writing:  chain(iterable, repeat(None)).
1459
1460   5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1461 */
1462
1463 static PyObject *
1464 imap_next(imapobject *lz)
1465 {
1466         PyObject *val;
1467         PyObject *argtuple;
1468         PyObject *result;
1469         int numargs, i;
1470
1471         numargs = PyTuple_Size(lz->iters);
1472         argtuple = PyTuple_New(numargs);
1473         if (argtuple == NULL)
1474                 return NULL;
1475
1476         for (i=0 ; i<numargs ; i++) {
1477                 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1478                 if (val == NULL) {
1479                         Py_DECREF(argtuple);
1480                         return NULL;
1481                 }
1482                 PyTuple_SET_ITEM(argtuple, i, val);
1483         }
1484         if (lz->func == Py_None)
1485                 return argtuple;
1486         result = PyObject_Call(lz->func, argtuple, NULL);
1487         Py_DECREF(argtuple);
1488         return result;
1489 }
1490
1491 PyDoc_STRVAR(imap_doc,
1492 "imap(func, *iterables) --> imap object\n\
1493 \n\
1494 Make an iterator that computes the function using arguments from\n\
1495 each of the iterables.  Like map() except that it returns\n\
1496 an iterator instead of a list and that it stops when the shortest\n\
1497 iterable is exhausted instead of filling in None for shorter\n\
1498 iterables.");
1499
1500 static PyTypeObject imap_type = {
1501         PyObject_HEAD_INIT(NULL)
1502         0,                              /* ob_size */
1503         "itertools.imap",               /* tp_name */
1504         sizeof(imapobject),             /* tp_basicsize */
1505         0,                              /* tp_itemsize */
1506         /* methods */
1507         (destructor)imap_dealloc,       /* tp_dealloc */
1508         0,                              /* tp_print */
1509         0,                              /* tp_getattr */
1510         0,                              /* tp_setattr */
1511         0,                              /* tp_compare */
1512         0,                              /* tp_repr */
1513         0,                              /* tp_as_number */
1514         0,                              /* tp_as_sequence */
1515         0,                              /* tp_as_mapping */
1516         0,                              /* tp_hash */
1517         0,                              /* tp_call */
1518         0,                              /* tp_str */
1519         PyObject_GenericGetAttr,        /* tp_getattro */
1520         0,                              /* tp_setattro */
1521         0,                              /* tp_as_buffer */
1522         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1523                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1524         imap_doc,                       /* tp_doc */
1525         (traverseproc)imap_traverse,    /* tp_traverse */
1526         0,                              /* tp_clear */
1527         0,                              /* tp_richcompare */
1528         0,                              /* tp_weaklistoffset */
1529         PyObject_SelfIter,              /* tp_iter */
1530         (iternextfunc)imap_next,        /* tp_iternext */
1531         0,                              /* tp_methods */
1532         0,                              /* tp_members */
1533         0,                              /* tp_getset */
1534         0,                              /* tp_base */
1535         0,                              /* tp_dict */
1536         0,                              /* tp_descr_get */
1537         0,                              /* tp_descr_set */
1538         0,                              /* tp_dictoffset */
1539         0,                              /* tp_init */
1540         0,                              /* tp_alloc */
1541         imap_new,                       /* tp_new */
1542         PyObject_GC_Del,                /* tp_free */
1543 };
1544
1545
1546 /* chain object ************************************************************/
1547
1548 typedef struct {
1549         PyObject_HEAD
1550         long    tuplesize;
1551         long    iternum;                /* which iterator is active */
1552         PyObject *ittuple;              /* tuple of iterators */
1553 } chainobject;
1554
1555 static PyTypeObject chain_type;
1556
1557 static PyObject *
1558 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1559 {
1560         chainobject *lz;
1561         int tuplesize = PySequence_Length(args);
1562         int i;
1563         PyObject *ittuple;
1564
1565         if (!_PyArg_NoKeywords("chain()", kwds))
1566                 return NULL;
1567
1568         /* obtain iterators */
1569         assert(PyTuple_Check(args));
1570         ittuple = PyTuple_New(tuplesize);
1571         if(ittuple == NULL)
1572                 return NULL;
1573         for (i=0; i < tuplesize; ++i) {
1574                 PyObject *item = PyTuple_GET_ITEM(args, i);
1575                 PyObject *it = PyObject_GetIter(item);
1576                 if (it == NULL) {
1577                         if (PyErr_ExceptionMatches(PyExc_TypeError))
1578                                 PyErr_Format(PyExc_TypeError,
1579                                     "chain argument #%d must support iteration",
1580                                     i+1);
1581                         Py_DECREF(ittuple);
1582                         return NULL;
1583                 }
1584                 PyTuple_SET_ITEM(ittuple, i, it);
1585         }
1586
1587         /* create chainobject structure */
1588         lz = (chainobject *)type->tp_alloc(type, 0);
1589         if (lz == NULL) {
1590                 Py_DECREF(ittuple);
1591                 return NULL;
1592         }
1593
1594         lz->ittuple = ittuple;
1595         lz->iternum = 0;
1596         lz->tuplesize = tuplesize;
1597
1598         return (PyObject *)lz;
1599 }
1600
1601 static void
1602 chain_dealloc(chainobject *lz)
1603 {
1604         PyObject_GC_UnTrack(lz);
1605         Py_XDECREF(lz->ittuple);
1606         lz->ob_type->tp_free(lz);
1607 }
1608
1609 static int
1610 chain_traverse(chainobject *lz, visitproc visit, void *arg)
1611 {
1612         Py_VISIT(lz->ittuple);
1613         return 0;
1614 }
1615
1616 static PyObject *
1617 chain_next(chainobject *lz)
1618 {
1619         PyObject *it;
1620         PyObject *item;
1621
1622         while (lz->iternum < lz->tuplesize) {
1623                 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1624                 item = PyIter_Next(it);
1625                 if (item != NULL)
1626                         return item;
1627                 if (PyErr_Occurred()) {
1628                         if (PyErr_ExceptionMatches(PyExc_StopIteration))
1629                                 PyErr_Clear();
1630                         else
1631                                 return NULL;
1632                 }
1633                 lz->iternum++;
1634         }
1635         return NULL;
1636 }
1637
1638 PyDoc_STRVAR(chain_doc,
1639 "chain(*iterables) --> chain object\n\
1640 \n\
1641 Return a chain object whose .next() method returns elements from the\n\
1642 first iterable until it is exhausted, then elements from the next\n\
1643 iterable, until all of the iterables are exhausted.");
1644
1645 static PyTypeObject chain_type = {
1646         PyObject_HEAD_INIT(NULL)
1647         0,                              /* ob_size */
1648         "itertools.chain",              /* tp_name */
1649         sizeof(chainobject),            /* tp_basicsize */
1650         0,                              /* tp_itemsize */
1651         /* methods */
1652         (destructor)chain_dealloc,      /* tp_dealloc */
1653         0,                              /* tp_print */
1654         0,                              /* tp_getattr */
1655         0,                              /* tp_setattr */
1656         0,                              /* tp_compare */
1657         0,                              /* tp_repr */
1658         0,                              /* tp_as_number */
1659         0,                              /* tp_as_sequence */
1660         0,                              /* tp_as_mapping */
1661         0,                              /* tp_hash */
1662         0,                              /* tp_call */
1663         0,                              /* tp_str */
1664         PyObject_GenericGetAttr,        /* tp_getattro */
1665         0,                              /* tp_setattro */
1666         0,                              /* tp_as_buffer */
1667         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1668                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1669         chain_doc,                      /* tp_doc */
1670         (traverseproc)chain_traverse,   /* tp_traverse */
1671         0,                              /* tp_clear */
1672         0,                              /* tp_richcompare */
1673         0,                              /* tp_weaklistoffset */
1674         PyObject_SelfIter,              /* tp_iter */
1675         (iternextfunc)chain_next,       /* tp_iternext */
1676         0,                              /* tp_methods */
1677         0,                              /* tp_members */
1678         0,                              /* tp_getset */
1679         0,                              /* tp_base */
1680         0,                              /* tp_dict */
1681         0,                              /* tp_descr_get */
1682         0,                              /* tp_descr_set */
1683         0,                              /* tp_dictoffset */
1684         0,                              /* tp_init */
1685         0,                              /* tp_alloc */
1686         chain_new,                      /* tp_new */
1687         PyObject_GC_Del,                /* tp_free */
1688 };
1689
1690
1691 /* ifilter object ************************************************************/
1692
1693 typedef struct {
1694         PyObject_HEAD
1695         PyObject *func;
1696         PyObject *it;
1697 } ifilterobject;
1698
1699 static PyTypeObject ifilter_type;
1700
1701 static PyObject *
1702 ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1703 {
1704         PyObject *func, *seq;
1705         PyObject *it;
1706         ifilterobject *lz;
1707
1708         if (!_PyArg_NoKeywords("ifilter()", kwds))
1709                 return NULL;
1710
1711         if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
1712                 return NULL;
1713
1714         /* Get iterator. */
1715         it = PyObject_GetIter(seq);
1716         if (it == NULL)
1717                 return NULL;
1718
1719         /* create ifilterobject structure */
1720         lz = (ifilterobject *)type->tp_alloc(type, 0);
1721         if (lz == NULL) {
1722                 Py_DECREF(it);
1723                 return NULL;
1724         }
1725         Py_INCREF(func);
1726         lz->func = func;
1727         lz->it = it;
1728
1729         return (PyObject *)lz;
1730 }
1731
1732 static void
1733 ifilter_dealloc(ifilterobject *lz)
1734 {
1735         PyObject_GC_UnTrack(lz);
1736         Py_XDECREF(lz->func);
1737         Py_XDECREF(lz->it);
1738         lz->ob_type->tp_free(lz);
1739 }
1740
1741 static int
1742 ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1743 {
1744         Py_VISIT(lz->it);
1745         Py_VISIT(lz->func);
1746         return 0;
1747 }
1748
1749 static PyObject *
1750 ifilter_next(ifilterobject *lz)
1751 {
1752         PyObject *item;
1753         PyObject *it = lz->it;
1754         long ok;
1755         PyObject *(*iternext)(PyObject *);
1756
1757         assert(PyIter_Check(it));
1758         iternext = *it->ob_type->tp_iternext;
1759         for (;;) {
1760                 item = iternext(it);
1761                 if (item == NULL)
1762                         return NULL;
1763
1764                 if (lz->func == Py_None) {
1765                         ok = PyObject_IsTrue(item);
1766                 } else {
1767                         PyObject *good;
1768                         good = PyObject_CallFunctionObjArgs(lz->func,
1769                                                             item, NULL);
1770                         if (good == NULL) {
1771                                 Py_DECREF(item);
1772                                 return NULL;
1773                         }
1774                         ok = PyObject_IsTrue(good);
1775                         Py_DECREF(good);
1776                 }
1777                 if (ok)
1778                         return item;
1779                 Py_DECREF(item);
1780         }
1781 }
1782
1783 PyDoc_STRVAR(ifilter_doc,
1784 "ifilter(function or None, sequence) --> ifilter object\n\
1785 \n\
1786 Return those items of sequence for which function(item) is true.\n\
1787 If function is None, return the items that are true.");
1788
1789 static PyTypeObject ifilter_type = {
1790         PyObject_HEAD_INIT(NULL)
1791         0,                              /* ob_size */
1792         "itertools.ifilter",            /* tp_name */
1793         sizeof(ifilterobject),          /* tp_basicsize */
1794         0,                              /* tp_itemsize */
1795         /* methods */
1796         (destructor)ifilter_dealloc,    /* tp_dealloc */
1797         0,                              /* tp_print */
1798         0,                              /* tp_getattr */
1799         0,                              /* tp_setattr */
1800         0,                              /* tp_compare */
1801         0,                              /* tp_repr */
1802         0,                              /* tp_as_number */
1803         0,                              /* tp_as_sequence */
1804         0,                              /* tp_as_mapping */
1805         0,                              /* tp_hash */
1806         0,                              /* tp_call */
1807         0,                              /* tp_str */
1808         PyObject_GenericGetAttr,        /* tp_getattro */
1809         0,                              /* tp_setattro */
1810         0,                              /* tp_as_buffer */
1811         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1812                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1813         ifilter_doc,                    /* tp_doc */
1814         (traverseproc)ifilter_traverse, /* tp_traverse */
1815         0,                              /* tp_clear */
1816         0,                              /* tp_richcompare */
1817         0,                              /* tp_weaklistoffset */
1818         PyObject_SelfIter,              /* tp_iter */
1819         (iternextfunc)ifilter_next,     /* tp_iternext */
1820         0,                              /* tp_methods */
1821         0,                              /* tp_members */
1822         0,                              /* tp_getset */
1823         0,                              /* tp_base */
1824         0,                              /* tp_dict */
1825         0,                              /* tp_descr_get */
1826         0,                              /* tp_descr_set */
1827         0,                              /* tp_dictoffset */
1828         0,                              /* tp_init */
1829         0,                              /* tp_alloc */
1830         ifilter_new,                    /* tp_new */
1831         PyObject_GC_Del,                /* tp_free */
1832 };
1833
1834
1835 /* ifilterfalse object ************************************************************/
1836
1837 typedef struct {
1838         PyObject_HEAD
1839         PyObject *func;
1840         PyObject *it;
1841 } ifilterfalseobject;
1842
1843 static PyTypeObject ifilterfalse_type;
1844
1845 static PyObject *
1846 ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1847 {
1848         PyObject *func, *seq;
1849         PyObject *it;
1850         ifilterfalseobject *lz;
1851
1852         if (!_PyArg_NoKeywords("ifilterfalse()", kwds))
1853                 return NULL;
1854
1855         if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1856                 return NULL;
1857
1858         /* Get iterator. */
1859         it = PyObject_GetIter(seq);
1860         if (it == NULL)
1861                 return NULL;
1862
1863         /* create ifilterfalseobject structure */
1864         lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1865         if (lz == NULL) {
1866                 Py_DECREF(it);
1867                 return NULL;
1868         }
1869         Py_INCREF(func);
1870         lz->func = func;
1871         lz->it = it;
1872
1873         return (PyObject *)lz;
1874 }
1875
1876 static void
1877 ifilterfalse_dealloc(ifilterfalseobject *lz)
1878 {
1879         PyObject_GC_UnTrack(lz);
1880         Py_XDECREF(lz->func);
1881         Py_XDECREF(lz->it);
1882         lz->ob_type->tp_free(lz);
1883 }
1884
1885 static int
1886 ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1887 {
1888         Py_VISIT(lz->it);
1889         Py_VISIT(lz->func);
1890         return 0;
1891 }
1892
1893 static PyObject *
1894 ifilterfalse_next(ifilterfalseobject *lz)
1895 {
1896         PyObject *item;
1897         PyObject *it = lz->it;
1898         long ok;
1899         PyObject *(*iternext)(PyObject *);
1900
1901         assert(PyIter_Check(it));
1902         iternext = *it->ob_type->tp_iternext;
1903         for (;;) {
1904                 item = iternext(it);
1905                 if (item == NULL)
1906                         return NULL;
1907
1908                 if (lz->func == Py_None) {
1909                         ok = PyObject_IsTrue(item);
1910                 } else {
1911                         PyObject *good;
1912                         good = PyObject_CallFunctionObjArgs(lz->func,
1913                                                             item, NULL);
1914                         if (good == NULL) {
1915                                 Py_DECREF(item);
1916                                 return NULL;
1917                         }
1918                         ok = PyObject_IsTrue(good);
1919                         Py_DECREF(good);
1920                 }
1921                 if (!ok)
1922                         return item;
1923                 Py_DECREF(item);
1924         }
1925 }
1926
1927 PyDoc_STRVAR(ifilterfalse_doc,
1928 "ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1929 \n\
1930 Return those items of sequence for which function(item) is false.\n\
1931 If function is None, return the items that are false.");
1932
1933 static PyTypeObject ifilterfalse_type = {
1934         PyObject_HEAD_INIT(NULL)
1935         0,                              /* ob_size */
1936         "itertools.ifilterfalse",       /* tp_name */
1937         sizeof(ifilterfalseobject),     /* tp_basicsize */
1938         0,                              /* tp_itemsize */
1939         /* methods */
1940         (destructor)ifilterfalse_dealloc,       /* tp_dealloc */
1941         0,                              /* tp_print */
1942         0,                              /* tp_getattr */
1943         0,                              /* tp_setattr */
1944         0,                              /* tp_compare */
1945         0,                              /* tp_repr */
1946         0,                              /* tp_as_number */
1947         0,                              /* tp_as_sequence */
1948         0,                              /* tp_as_mapping */
1949         0,                              /* tp_hash */
1950         0,                              /* tp_call */
1951         0,                              /* tp_str */
1952         PyObject_GenericGetAttr,        /* tp_getattro */
1953         0,                              /* tp_setattro */
1954         0,                              /* tp_as_buffer */
1955         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1956                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
1957         ifilterfalse_doc,               /* tp_doc */
1958         (traverseproc)ifilterfalse_traverse,    /* tp_traverse */
1959         0,                              /* tp_clear */
1960         0,                              /* tp_richcompare */
1961         0,                              /* tp_weaklistoffset */
1962         PyObject_SelfIter,              /* tp_iter */
1963         (iternextfunc)ifilterfalse_next,        /* tp_iternext */
1964         0,                              /* tp_methods */
1965         0,                              /* tp_members */
1966         0,                              /* tp_getset */
1967         0,                              /* tp_base */
1968         0,                              /* tp_dict */
1969         0,                              /* tp_descr_get */
1970         0,                              /* tp_descr_set */
1971         0,                              /* tp_dictoffset */
1972         0,                              /* tp_init */
1973         0,                              /* tp_alloc */
1974         ifilterfalse_new,               /* tp_new */
1975         PyObject_GC_Del,                /* tp_free */
1976 };
1977
1978
1979 /* count object ************************************************************/
1980
1981 typedef struct {
1982         PyObject_HEAD
1983         long    cnt;
1984 } countobject;
1985
1986 static PyTypeObject count_type;
1987
1988 static PyObject *
1989 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1990 {
1991         countobject *lz;
1992         long cnt = 0;
1993
1994         if (!_PyArg_NoKeywords("count()", kwds))
1995                 return NULL;
1996
1997         if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1998                 return NULL;
1999
2000         /* create countobject structure */
2001         lz = (countobject *)PyObject_New(countobject, &count_type);
2002         if (lz == NULL)
2003                 return NULL;
2004         lz->cnt = cnt;
2005
2006         return (PyObject *)lz;
2007 }
2008
2009 static PyObject *
2010 count_next(countobject *lz)
2011 {
2012         return PyInt_FromLong(lz->cnt++);
2013 }
2014
2015 static PyObject *
2016 count_repr(countobject *lz)
2017 {
2018         return PyString_FromFormat("count(%ld)", lz->cnt);
2019 }
2020
2021 PyDoc_STRVAR(count_doc,
2022 "count([firstval]) --> count object\n\
2023 \n\
2024 Return a count object whose .next() method returns consecutive\n\
2025 integers starting from zero or, if specified, from firstval.");
2026
2027 static PyTypeObject count_type = {
2028         PyObject_HEAD_INIT(NULL)
2029         0,                              /* ob_size */
2030         "itertools.count",              /* tp_name */
2031         sizeof(countobject),            /* tp_basicsize */
2032         0,                              /* tp_itemsize */
2033         /* methods */
2034         (destructor)PyObject_Del,       /* tp_dealloc */
2035         0,                              /* tp_print */
2036         0,                              /* tp_getattr */
2037         0,                              /* tp_setattr */
2038         0,                              /* tp_compare */
2039         (reprfunc)count_repr,           /* tp_repr */
2040         0,                              /* tp_as_number */
2041         0,                              /* tp_as_sequence */
2042         0,                              /* tp_as_mapping */
2043         0,                              /* tp_hash */
2044         0,                              /* tp_call */
2045         0,                              /* tp_str */
2046         PyObject_GenericGetAttr,        /* tp_getattro */
2047         0,                              /* tp_setattro */
2048         0,                              /* tp_as_buffer */
2049         Py_TPFLAGS_DEFAULT,             /* tp_flags */
2050         count_doc,                      /* tp_doc */
2051         0,                              /* tp_traverse */
2052         0,                              /* tp_clear */
2053         0,                              /* tp_richcompare */
2054         0,                              /* tp_weaklistoffset */
2055         PyObject_SelfIter,              /* tp_iter */
2056         (iternextfunc)count_next,       /* tp_iternext */
2057         0,                              /* tp_methods */
2058         0,                              /* tp_members */
2059         0,                              /* tp_getset */
2060         0,                              /* tp_base */
2061         0,                              /* tp_dict */
2062         0,                              /* tp_descr_get */
2063         0,                              /* tp_descr_set */
2064         0,                              /* tp_dictoffset */
2065         0,                              /* tp_init */
2066         0,                              /* tp_alloc */
2067         count_new,                      /* tp_new */
2068 };
2069
2070
2071 /* izip object ************************************************************/
2072
2073 #include "Python.h"
2074
2075 typedef struct {
2076         PyObject_HEAD
2077         long    tuplesize;
2078         PyObject *ittuple;              /* tuple of iterators */
2079         PyObject *result;
2080 } izipobject;
2081
2082 static PyTypeObject izip_type;
2083
2084 static PyObject *
2085 izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2086 {
2087         izipobject *lz;
2088         int i;
2089         PyObject *ittuple;  /* tuple of iterators */
2090         PyObject *result;
2091         int tuplesize = PySequence_Length(args);
2092
2093         if (!_PyArg_NoKeywords("izip()", kwds))
2094                 return NULL;
2095
2096         /* args must be a tuple */
2097         assert(PyTuple_Check(args));
2098
2099         /* obtain iterators */
2100         ittuple = PyTuple_New(tuplesize);
2101         if(ittuple == NULL)
2102                 return NULL;
2103         for (i=0; i < tuplesize; ++i) {
2104                 PyObject *item = PyTuple_GET_ITEM(args, i);
2105                 PyObject *it = PyObject_GetIter(item);
2106                 if (it == NULL) {
2107                         if (PyErr_ExceptionMatches(PyExc_TypeError))
2108                                 PyErr_Format(PyExc_TypeError,
2109                                     "izip argument #%d must support iteration",
2110                                     i+1);
2111                         Py_DECREF(ittuple);
2112                         return NULL;
2113                 }
2114                 PyTuple_SET_ITEM(ittuple, i, it);
2115         }
2116
2117         /* create a result holder */
2118         result = PyTuple_New(tuplesize);
2119         if (result == NULL) {
2120                 Py_DECREF(ittuple);
2121                 return NULL;
2122         }
2123         for (i=0 ; i < tuplesize ; i++) {
2124                 Py_INCREF(Py_None);
2125                 PyTuple_SET_ITEM(result, i, Py_None);
2126         }
2127
2128         /* create izipobject structure */
2129         lz = (izipobject *)type->tp_alloc(type, 0);
2130         if (lz == NULL) {
2131                 Py_DECREF(ittuple);
2132                 Py_DECREF(result);
2133                 return NULL;
2134         }
2135         lz->ittuple = ittuple;
2136         lz->tuplesize = tuplesize;
2137         lz->result = result;
2138
2139         return (PyObject *)lz;
2140 }
2141
2142 static void
2143 izip_dealloc(izipobject *lz)
2144 {
2145         PyObject_GC_UnTrack(lz);
2146         Py_XDECREF(lz->ittuple);
2147         Py_XDECREF(lz->result);
2148         lz->ob_type->tp_free(lz);
2149 }
2150
2151 static int
2152 izip_traverse(izipobject *lz, visitproc visit, void *arg)
2153 {
2154         Py_VISIT(lz->ittuple);
2155         Py_VISIT(lz->result);
2156         return 0;
2157 }
2158
2159 static PyObject *
2160 izip_next(izipobject *lz)
2161 {
2162         int i;
2163         long tuplesize = lz->tuplesize;
2164         PyObject *result = lz->result;
2165         PyObject *it;
2166         PyObject *item;
2167         PyObject *olditem;
2168
2169         if (tuplesize == 0)
2170                 return NULL;
2171         if (result->ob_refcnt == 1) {
2172                 Py_INCREF(result);
2173                 for (i=0 ; i < tuplesize ; i++) {
2174                         it = PyTuple_GET_ITEM(lz->ittuple, i);
2175                         assert(PyIter_Check(it));
2176                         item = (*it->ob_type->tp_iternext)(it);
2177                         if (item == NULL) {
2178                                 Py_DECREF(result);
2179                                 return NULL;
2180                         }
2181                         olditem = PyTuple_GET_ITEM(result, i);
2182                         PyTuple_SET_ITEM(result, i, item);
2183                         Py_DECREF(olditem);
2184                 }
2185         } else {
2186                 result = PyTuple_New(tuplesize);
2187                 if (result == NULL)
2188                         return NULL;
2189                 for (i=0 ; i < tuplesize ; i++) {
2190                         it = PyTuple_GET_ITEM(lz->ittuple, i);
2191                         assert(PyIter_Check(it));
2192                         item = (*it->ob_type->tp_iternext)(it);
2193                         if (item == NULL) {
2194                                 Py_DECREF(result);
2195                                 return NULL;
2196                         }
2197                         PyTuple_SET_ITEM(result, i, item);
2198                 }
2199         }
2200         return result;
2201 }
2202
2203 PyDoc_STRVAR(izip_doc,
2204 "izip(iter1 [,iter2 [...]]) --> izip object\n\
2205 \n\
2206 Return a izip object whose .next() method returns a tuple where\n\
2207 the i-th element comes from the i-th iterable argument.  The .next()\n\
2208 method continues until the shortest iterable in the argument sequence\n\
2209 is exhausted and then it raises StopIteration.  Works like the zip()\n\
2210 function but consumes less memory by returning an iterator instead of\n\
2211 a list.");
2212
2213 static PyTypeObject izip_type = {
2214         PyObject_HEAD_INIT(NULL)
2215         0,                              /* ob_size */
2216         "itertools.izip",               /* tp_name */
2217         sizeof(izipobject),             /* tp_basicsize */
2218         0,                              /* tp_itemsize */
2219         /* methods */
2220         (destructor)izip_dealloc,       /* tp_dealloc */
2221         0,                              /* tp_print */
2222         0,                              /* tp_getattr */
2223         0,                              /* tp_setattr */
2224         0,                              /* tp_compare */
2225         0,                              /* tp_repr */
2226         0,                              /* tp_as_number */
2227         0,                              /* tp_as_sequence */
2228         0,                              /* tp_as_mapping */
2229         0,                              /* tp_hash */
2230         0,                              /* tp_call */
2231         0,                              /* tp_str */
2232         PyObject_GenericGetAttr,        /* tp_getattro */
2233         0,                              /* tp_setattro */
2234         0,                              /* tp_as_buffer */
2235         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2236                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
2237         izip_doc,                       /* tp_doc */
2238         (traverseproc)izip_traverse,    /* tp_traverse */
2239         0,                              /* tp_clear */
2240         0,                              /* tp_richcompare */
2241         0,                              /* tp_weaklistoffset */
2242         PyObject_SelfIter,              /* tp_iter */
2243         (iternextfunc)izip_next,        /* tp_iternext */
2244         0,                              /* tp_methods */
2245         0,                              /* tp_members */
2246         0,                              /* tp_getset */
2247         0,                              /* tp_base */
2248         0,                              /* tp_dict */
2249         0,                              /* tp_descr_get */
2250         0,                              /* tp_descr_set */
2251         0,                              /* tp_dictoffset */
2252         0,                              /* tp_init */
2253         0,                              /* tp_alloc */
2254         izip_new,                       /* tp_new */
2255         PyObject_GC_Del,                /* tp_free */
2256 };
2257
2258
2259 /* repeat object ************************************************************/
2260
2261 typedef struct {
2262         PyObject_HEAD
2263         PyObject *element;
2264         long cnt;
2265 } repeatobject;
2266
2267 static PyTypeObject repeat_type;
2268
2269 static PyObject *
2270 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2271 {
2272         repeatobject *ro;
2273         PyObject *element;
2274         long cnt = -1;
2275
2276         if (!_PyArg_NoKeywords("repeat()", kwds))
2277                 return NULL;
2278
2279         if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
2280                 return NULL;
2281
2282         if (PyTuple_Size(args) == 2 && cnt < 0)
2283                 cnt = 0;
2284
2285         ro = (repeatobject *)type->tp_alloc(type, 0);
2286         if (ro == NULL)
2287                 return NULL;
2288         Py_INCREF(element);
2289         ro->element = element;
2290         ro->cnt = cnt;
2291         return (PyObject *)ro;
2292 }
2293
2294 static void
2295 repeat_dealloc(repeatobject *ro)
2296 {
2297         PyObject_GC_UnTrack(ro);
2298         Py_XDECREF(ro->element);
2299         ro->ob_type->tp_free(ro);
2300 }
2301
2302 static int
2303 repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2304 {
2305         Py_VISIT(ro->element);
2306         return 0;
2307 }
2308
2309 static PyObject *
2310 repeat_next(repeatobject *ro)
2311 {
2312         if (ro->cnt == 0)
2313                 return NULL;
2314         if (ro->cnt > 0)
2315                 ro->cnt--;
2316         Py_INCREF(ro->element);
2317         return ro->element;
2318 }
2319
2320 static PyObject *
2321 repeat_repr(repeatobject *ro)
2322 {
2323         PyObject *result, *objrepr;
2324
2325         objrepr = PyObject_Repr(ro->element);
2326         if (objrepr == NULL)
2327                 return NULL;
2328
2329         if (ro->cnt == -1)
2330                 result = PyString_FromFormat("repeat(%s)",
2331                         PyString_AS_STRING(objrepr));
2332         else
2333                 result = PyString_FromFormat("repeat(%s, %ld)",
2334                         PyString_AS_STRING(objrepr), ro->cnt);
2335         Py_DECREF(objrepr);
2336         return result;
2337 }
2338
2339 static PyObject *
2340 repeat_len(repeatobject *ro)
2341 {
2342         if (ro->cnt == -1) {
2343                 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2344                 return NULL;
2345         }
2346         return PyInt_FromLong(ro->cnt);
2347 }
2348
2349 PyDoc_STRVAR(length_cue_doc, "Private method returning an estimate of len(list(it)).");
2350
2351 static PyMethodDef repeat_methods[] = {
2352         {"_length_cue", (PyCFunction)repeat_len, METH_NOARGS, length_cue_doc},
2353         {NULL,          NULL}           /* sentinel */
2354 };
2355
2356 PyDoc_STRVAR(repeat_doc,
2357 "repeat(element [,times]) -> create an iterator which returns the element\n\
2358 for the specified number of times.  If not specified, returns the element\n\
2359 endlessly.");
2360
2361 static PyTypeObject repeat_type = {
2362         PyObject_HEAD_INIT(NULL)
2363         0,                              /* ob_size */
2364         "itertools.repeat",             /* tp_name */
2365         sizeof(repeatobject),           /* tp_basicsize */
2366         0,                              /* tp_itemsize */
2367         /* methods */
2368         (destructor)repeat_dealloc,     /* tp_dealloc */
2369         0,                              /* tp_print */
2370         0,                              /* tp_getattr */
2371         0,                              /* tp_setattr */
2372         0,                              /* tp_compare */
2373         (reprfunc)repeat_repr,          /* tp_repr */
2374         0,                              /* tp_as_number */
2375         0,                              /* tp_as_sequence */
2376         0,                              /* tp_as_mapping */
2377         0,                              /* tp_hash */
2378         0,                              /* tp_call */
2379         0,                              /* tp_str */
2380         PyObject_GenericGetAttr,        /* tp_getattro */
2381         0,                              /* tp_setattro */
2382         0,                              /* tp_as_buffer */
2383         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2384                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
2385         repeat_doc,                     /* tp_doc */
2386         (traverseproc)repeat_traverse,  /* tp_traverse */
2387         0,                              /* tp_clear */
2388         0,                              /* tp_richcompare */
2389         0,                              /* tp_weaklistoffset */
2390         PyObject_SelfIter,              /* tp_iter */
2391         (iternextfunc)repeat_next,      /* tp_iternext */
2392         repeat_methods,                 /* tp_methods */
2393         0,                              /* tp_members */
2394         0,                              /* tp_getset */
2395         0,                              /* tp_base */
2396         0,                              /* tp_dict */
2397         0,                              /* tp_descr_get */
2398         0,                              /* tp_descr_set */
2399         0,                              /* tp_dictoffset */
2400         0,                              /* tp_init */
2401         0,                              /* tp_alloc */
2402         repeat_new,                     /* tp_new */
2403         PyObject_GC_Del,                /* tp_free */
2404 };
2405
2406
2407 /* module level code ********************************************************/
2408
2409 PyDoc_STRVAR(module_doc,
2410 "Functional tools for creating and using iterators.\n\
2411 \n\
2412 Infinite iterators:\n\
2413 count([n]) --> n, n+1, n+2, ...\n\
2414 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
2415 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
2416 \n\
2417 Iterators terminating on the shortest input sequence:\n\
2418 izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
2419 ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2420 ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
2421 islice(seq, [start,] stop [, step]) --> elements from\n\
2422        seq[start:stop:step]\n\
2423 imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2424 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
2425 tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
2426 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
2427 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2428 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
2429 groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
2430 ");
2431
2432
2433 static PyMethodDef module_methods[] = {
2434         {"tee", (PyCFunction)tee,       METH_VARARGS, tee_doc},
2435         {NULL,          NULL}           /* sentinel */
2436 };
2437
2438 PyMODINIT_FUNC
2439 inititertools(void)
2440 {
2441         int i;
2442         PyObject *m;
2443         char *name;
2444         PyTypeObject *typelist[] = {
2445                 &cycle_type,
2446                 &dropwhile_type,
2447                 &takewhile_type,
2448                 &islice_type,
2449                 &starmap_type,
2450                 &imap_type,
2451                 &chain_type,
2452                 &ifilter_type,
2453                 &ifilterfalse_type,
2454                 &count_type,
2455                 &izip_type,
2456                 &repeat_type,
2457                 &groupby_type,
2458                 NULL
2459         };
2460
2461         teedataobject_type.ob_type = &PyType_Type;
2462         m = Py_InitModule3("itertools", module_methods, module_doc);
2463
2464         for (i=0 ; typelist[i] != NULL ; i++) {
2465                 if (PyType_Ready(typelist[i]) < 0)
2466                         return;
2467                 name = strchr(typelist[i]->tp_name, '.');
2468                 assert (name != NULL);
2469                 Py_INCREF(typelist[i]);
2470                 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
2471         }
2472
2473         if (PyType_Ready(&teedataobject_type) < 0)
2474                 return;
2475         if (PyType_Ready(&tee_type) < 0)
2476                 return;
2477         if (PyType_Ready(&_grouper_type) < 0)
2478                 return;
2479 }