#define PY_SSIZE_T_CLEAN

#include "Python.h"

#include "structmember.h"

typedef struct {

PyObject_HEAD

PyObject *it;

PyObject *keyfunc;

PyObject *tgtkey;

PyObject *currkey;

PyObject *currvalue;

} groupbyobject;

static PyTypeObject groupby_type;

static PyObject *_grouper_create(groupbyobject *, PyObject *);

static PyObject *

groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)

{

static char *kwargs[] = {"iterable", "key", NULL};

groupbyobject *gbo;

PyObject *it, *keyfunc = Py_None;

if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,

&it, &keyfunc))

return NULL;

gbo = (groupbyobject *)type->tp_alloc(type, 0);

if (gbo == NULL)

return NULL;

gbo->tgtkey = NULL;

gbo->currkey = NULL;

gbo->currvalue = NULL;

gbo->keyfunc = keyfunc;

Py_INCREF(keyfunc);

gbo->it = PyObject_GetIter(it);

if (gbo->it == NULL) {

Py_DECREF(gbo);

return NULL;

}

return (PyObject *)gbo;

}

static void

groupby_dealloc(groupbyobject *gbo)

{

PyObject_GC_UnTrack(gbo);

Py_XDECREF(gbo->it);

Py_XDECREF(gbo->keyfunc);

Py_XDECREF(gbo->tgtkey);

Py_XDECREF(gbo->currkey);

Py_XDECREF(gbo->currvalue);

Py_TYPE(gbo)->tp_free(gbo);

}

static int

groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)

{

Py_VISIT(gbo->it);

Py_VISIT(gbo->keyfunc);

Py_VISIT(gbo->tgtkey);

Py_VISIT(gbo->currkey);

Py_VISIT(gbo->currvalue);

return 0;

}

static PyObject *

groupby_next(groupbyobject *gbo)

{

PyObject *newvalue, *newkey, *r, *grouper;

/* skip to next iteration group */

for (;;) {

if (gbo->currkey == NULL)

/* pass */;

else if (gbo->tgtkey == NULL)

break;

else {

int rcmp;

rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);

if (rcmp == -1)

return NULL;

else if (rcmp == 0)

break;

}

newvalue = PyIter_Next(gbo->it);

if (newvalue == NULL)

return NULL;

if (gbo->keyfunc == Py_None) {

newkey = newvalue;

Py_INCREF(newvalue);

} else {

newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);

if (newkey == NULL) {

Py_DECREF(newvalue);

return NULL;

}

}

Py_XSETREF(gbo->currkey, newkey);

Py_XSETREF(gbo->currvalue, newvalue);

}

Py_INCREF(gbo->currkey);

Py_XSETREF(gbo->tgtkey, gbo->currkey);

grouper = _grouper_create(gbo, gbo->tgtkey);

if (grouper == NULL)

return NULL;

r = PyTuple_Pack(2, gbo->currkey, grouper);

Py_DECREF(grouper);

return r;

}

static PyObject *

groupby_reduce(groupbyobject *lz)

{

/* reduce as a 'new' call with an optional 'setstate' if groupby

PyObject *value;

if (lz->tgtkey && lz->currkey && lz->currvalue)

value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),

lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);

else

value = Py_BuildValue("O(OO)", Py_TYPE(lz),

lz->it, lz->keyfunc);

return value;

}

PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");

static PyObject *

groupby_setstate(groupbyobject *lz, PyObject *state)

{

PyObject *currkey, *currvalue, *tgtkey;

if (!PyTuple_Check(state)) {

PyErr_SetString(PyExc_TypeError, "state is not a tuple");

return NULL;

}

if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {

return NULL;

}

Py_INCREF(currkey);

Py_XSETREF(lz->currkey, currkey);

Py_INCREF(currvalue);

Py_XSETREF(lz->currvalue, currvalue);

Py_INCREF(tgtkey);

Py_XSETREF(lz->tgtkey, tgtkey);

Py_RETURN_NONE;

}

PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");

static PyMethodDef groupby_methods[] = {

{"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,

reduce_doc},

{"__setstate__", (PyCFunction)groupby_setstate, METH_O,

setstate_doc},

{NULL, NULL} /* sentinel */

};

PyDoc_STRVAR(groupby_doc,

(key, sub-iterator) grouped by each value of key(value).\n");

static PyTypeObject groupby_type = {

PyVarObject_HEAD_INIT(NULL, 0)

"itertools.groupby", /* tp_name */

sizeof(groupbyobject), /* tp_basicsize */

0, /* tp_itemsize */

/* methods */

(destructor)groupby_dealloc, /* tp_dealloc */

0, /* tp_print */

0, /* tp_getattr */

0, /* tp_setattr */

0, /* tp_reserved */

0, /* tp_repr */

0, /* tp_as_number */

0, /* tp_as_sequence */

0, /* tp_as_mapping */

0, /* tp_hash */

0, /* tp_call */

0, /* tp_str */

PyObject_GenericGetAttr, /* tp_getattro */

0, /* tp_setattro */

0, /* tp_as_buffer */

Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |

Py_TPFLAGS_BASETYPE, /* tp_flags */

groupby_doc, /* tp_doc */

(traverseproc)groupby_traverse, /* tp_traverse */

0, /* tp_clear */

0, /* tp_richcompare */

0, /* tp_weaklistoffset */

PyObject_SelfIter, /* tp_iter */

(iternextfunc)groupby_next, /* tp_iternext */

groupby_methods, /* tp_methods */

0, /* tp_members */

0, /* tp_getset */

0, /* tp_base */

0, /* tp_dict */

0, /* tp_descr_get */

0, /* tp_descr_set */

0, /* tp_dictoffset */

0, /* tp_init */

0, /* tp_alloc */

groupby_new, /* tp_new */

PyObject_GC_Del, /* tp_free */

};

typedef struct {

PyObject_HEAD

PyObject *parent;

PyObject *tgtkey;

} _grouperobject;

static PyTypeObject _grouper_type;

static PyObject *

_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)

{

PyObject *parent, *tgtkey;

if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))

return NULL;

return _grouper_create((groupbyobject*) parent, tgtkey);

}

static PyObject *

_grouper_create(groupbyobject *parent, PyObject *tgtkey)

{

_grouperobject *igo;

igo = PyObject_GC_New(_grouperobject, &_grouper_type);

if (igo == NULL)

return NULL;

igo->parent = (PyObject *)parent;

Py_INCREF(parent);

igo->tgtkey = tgtkey;

Py_INCREF(tgtkey);

PyObject_GC_Track(igo);

return (PyObject *)igo;

}

static void

_grouper_dealloc(_grouperobject *igo)

{

PyObject_GC_UnTrack(igo);

Py_DECREF(igo->parent);

Py_DECREF(igo->tgtkey);

PyObject_GC_Del(igo);

}

static int

_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)

{

Py_VISIT(igo->parent);

Py_VISIT(igo->tgtkey);

return 0;

}

static PyObject *

_grouper_next(_grouperobject *igo)

{

groupbyobject *gbo = (groupbyobject *)igo->parent;

PyObject *newvalue, *newkey, *r;

int rcmp;

if (gbo->currvalue == NULL) {

newvalue = PyIter_Next(gbo->it);

if (newvalue == NULL)

return NULL;

if (gbo->keyfunc == Py_None) {

newkey = newvalue;

Py_INCREF(newvalue);

} else {

newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);

if (newkey == NULL) {

Py_DECREF(newvalue);

return NULL;

}

}

assert(gbo->currkey == NULL);

gbo->currkey = newkey;

gbo->currvalue = newvalue;

}

assert(gbo->currkey != NULL);

rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);

if (rcmp <= 0)

/* got any error or current group is end */

return NULL;

r = gbo->currvalue;

gbo->currvalue = NULL;

Py_CLEAR(gbo->currkey);

return r;

}

static PyObject *

_grouper_reduce(_grouperobject *lz)

{

return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);

}

static PyMethodDef _grouper_methods[] = {

{"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,

reduce_doc},

{NULL, NULL} /* sentinel */

};

static PyTypeObject _grouper_type = {

PyVarObject_HEAD_INIT(NULL, 0)

"itertools._grouper", /* tp_name */

sizeof(_grouperobject), /* tp_basicsize */

0, /* tp_itemsize */

/* methods */

(destructor)_grouper_dealloc, /* tp_dealloc */

0, /* tp_print */

0, /* tp_getattr */

0, /* tp_setattr */

0, /* tp_reserved */

0, /* tp_repr */

0, /* tp_as_number */

0, /* tp_as_sequence */

0, /* tp_as_mapping */

0, /* tp_hash */

0, /* tp_call */

0, /* tp_str */

PyObject_GenericGetAttr, /* tp_getattro */

0, /* tp_setattro */

0, /* tp_as_buffer */

Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */

0, /* tp_doc */

(traverseproc)_grouper_traverse, /* tp_traverse */

0, /* tp_clear */

0, /* tp_richcompare */

0, /* tp_weaklistoffset */

PyObject_SelfIter, /* tp_iter */

(iternextfunc)_grouper_next, /* tp_iternext */

_grouper_methods, /* tp_methods */

0, /* tp_members */

0, /* tp_getset */

0, /* tp_base */

0, /* tp_dict */

0, /* tp_descr_get */

0, /* tp_descr_set */

0, /* tp_dictoffset */

0, /* tp_init */

0, /* tp_alloc */

_grouper_new, /* tp_new */

PyObject_GC_Del, /* tp_free */

};

#define LINKCELLS 57

typedef struct {

PyObject_HEAD

PyObject *it;

int numread; /* 0 <= numread <= LINKCELLS */

PyObject *nextlink;

PyObject *(values[LINKCELLS]);

} teedataobject;

typedef struct {

PyObject_HEAD

teedataobject *dataobj;

int index; /* 0 <= index <= LINKCELLS */

PyObject *weakreflist;

} teeobject;

static PyTypeObject teedataobject_type;

static PyObject *

teedataobject_newinternal(PyObject *it)

{

teedataobject *tdo;

tdo = PyObject_GC_New(teedataobject, &teedataobject_type);

if (tdo == NULL)

return NULL;

tdo->numread = 0;

tdo->nextlink = NULL;

Py_INCREF(it);

tdo->it = it;

PyObject_GC_Track(tdo);

return (PyObject *)tdo;

}

static PyObject *

teedataobject_jumplink(teedataobject *tdo)

{

if (tdo->nextlink == NULL)

tdo->nextlink = teedataobject_newinternal(tdo->it);

Py_XINCREF(tdo->nextlink);

return tdo->nextlink;

}

static PyObject *

teedataobject_getitem(teedataobject *tdo, int i)

{

PyObject *value;

assert(i < LINKCELLS);

if (i < tdo->numread)

value = tdo->values[i];

else {

/* this is the lead iterator, so fetch more data */

assert(i == tdo->numread);

value = PyIter_Next(tdo->it);

if (value == NULL)

return NULL;

tdo->numread++;

tdo->values[i] = value;

}

Py_INCREF(value);

return value;

}

static int

teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)

{

int i;

Py_VISIT(tdo->it);

for (i = 0; i < tdo->numread; i++)

Py_VISIT(tdo->values[i]);

Py_VISIT(tdo->nextlink);

return 0;

}

static void

teedataobject_safe_decref(PyObject *obj)

{

while (obj && Py_TYPE(obj) == &teedataobject_type &&

Py_REFCNT(obj) == 1) {

PyObject *nextlink = ((teedataobject *)obj)->nextlink;

((teedataobject *)obj)->nextlink = NULL;

Py_DECREF(obj);

obj = nextlink;

}

Py_XDECREF(obj);

}

static int

teedataobject_clear(teedataobject *tdo)

{

int i;

PyObject *tmp;

Py_CLEAR(tdo->it);

for (i=0 ; i<tdo->numread ; i++)

Py_CLEAR(tdo->values[i]);

tmp = tdo->nextlink;

tdo->nextlink = NULL;

teedataobject_safe_decref(tmp);

return 0;

}

static void

teedataobject_dealloc(teedataobject *tdo)

{

PyObject_GC_UnTrack(tdo);

teedataobject_clear(tdo);

PyObject_GC_Del(tdo);

}

static PyObject *

teedataobject_reduce(teedataobject *tdo)

{

int i;

/* create a temporary list of already iterated values */

PyObject *values = PyList_New(tdo->numread);

if (!values)

return NULL;

for (i=0 ; i<tdo->numread ; i++) {

Py_INCREF(tdo->values[i]);

PyList_SET_ITEM(values, i, tdo->values[i]);

}

return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,

values,

tdo->nextlink ? tdo->nextlink : Py_None);

}

static PyTypeObject teedataobject_type;

static PyObject *

teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)

{

teedataobject *tdo;

PyObject *it, *values, *next;

Py_ssize_t i, len;

assert(type == &teedataobject_type);

if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))

return NULL;

tdo = (teedataobject *)teedataobject_newinternal(it);

if (!tdo)

return NULL;

len = PyList_GET_SIZE(values);

if (len > LINKCELLS)

goto err;

for (i=0; i<len; i++) {

tdo->values[i] = PyList_GET_ITEM(values, i);

Py_INCREF(tdo->values[i]);

}

/* len <= LINKCELLS < INT_MAX */

tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);

if (len == LINKCELLS) {

if (next != Py_None) {

if (Py_TYPE(next) != &teedataobject_type)

goto err;

assert(tdo->nextlink == NULL);

Py_INCREF(next);

tdo->nextlink = next;

}

} else {

if (next != Py_None)

goto err; /* shouldn't have a next if we are not full */

}

return (PyObject*)tdo;

err:

Py_XDECREF(tdo);

PyErr_SetString(PyExc_ValueError, "Invalid arguments");

return NULL;

}

static PyMethodDef teedataobject_methods[] = {

{"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,

reduce_doc},

{NULL, NULL} /* sentinel */

};

PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");

static PyTypeObject teedataobject_type = {

PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */

"itertools._tee_dataobject", /* tp_name */

sizeof(teedataobject), /* tp_basicsize */

0, /* tp_itemsize */

/* methods */

(destructor)teedataobject_dealloc, /* tp_dealloc */

0, /* tp_print */

0, /* tp_getattr */

0, /* tp_setattr */

0, /* tp_reserved */

0, /* tp_repr */

0, /* tp_as_number */

0, /* tp_as_sequence */

0, /* tp_as_mapping */

0, /* tp_hash */

0, /* tp_call */

0, /* tp_str */

PyObject_GenericGetAttr, /* tp_getattro */

0, /* tp_setattro */

0, /* tp_as_buffer */

Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */

teedataobject_doc, /* tp_doc */

(traverseproc)teedataobject_traverse, /* tp_traverse */

(inquiry)teedataobject_clear, /* tp_clear */

0, /* tp_richcompare */

0, /* tp_weaklistoffset */

0, /* tp_iter */

0, /* tp_iternext */

teedataobject_methods, /* tp_methods */

0, /* tp_members */

0, /* tp_getset */

0, /* tp_base */

0, /* tp_dict */

0, /* tp_descr_get */

0, /* tp_descr_set */

0, /* tp_dictoffset */

0, /* tp_init */

0, /* tp_alloc */

teedataobject_new, /* tp_new */

PyObject_GC_Del, /* tp_free */

};

static PyTypeObject tee_type;

static PyObject *

tee_next(teeobject *to)

{

PyObject *value, *link;

if (to->index >= LINKCELLS) {

link = teedataobject_jumplink(to->dataobj);

if (link == NULL)

return NULL;

Py_SETREF(to->dataobj, (teedataobject *)link);

to->index = 0;

}

value = teedataobject_getitem(to->dataobj, to->index);

if (value == NULL)

return NULL;

to->index++;

return value;

}

static int

tee_traverse(teeobject *to, visitproc visit, void *arg)

{

Py_VISIT((PyObject *)to->dataobj);

return 0;

}

static PyObject *

tee_copy(teeobject *to)

{

teeobject *newto;

newto = PyObject_GC_New(teeobject, &tee_type);

if (newto == NULL)

return NULL;

Py_INCREF(to->dataobj);

newto->dataobj = to->dataobj;

newto->index = to->index;

newto->weakreflist = NULL;

PyObject_GC_Track(newto);

return (PyObject *)newto;

}

PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");

static PyObject *

tee_fromiterable(PyObject *iterable)

{

teeobject *to;

PyObject *it = NULL;

it = PyObject_GetIter(iterable);

if (it == NULL)

return NULL;

if (PyObject_TypeCheck(it, &tee_type)) {

to = (teeobject *)tee_copy((teeobject *)it);

goto done;

}

to = PyObject_GC_New(teeobject, &tee_type);

if (to == NULL)

goto done;

to->dataobj = (teedataobject *)teedataobject_newinternal(it);

if (!to->dataobj) {

PyObject_GC_Del(to);

to = NULL;

goto done;

}

to->index = 0;

to->weakreflist = NULL;

PyObject_GC_Track(to);

done:

Py_XDECREF(it);

return (PyObject *)to;

}

static PyObject *

tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)

{

PyObject *iterable;

if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))

return NULL;

return tee_fromiterable(iterable);

}

static int

tee_clear(teeobject *to)

{

if (to->weakreflist != NULL)

PyObject_ClearWeakRefs((PyObject *) to);

Py_CLEAR(to->dataobj);

return 0;

}

static void

tee_dealloc(teeobject *to)

{

PyObject_GC_UnTrack(to);

tee_clear(to);

PyObject_GC_Del(to);

}

static PyObject *

tee_reduce(teeobject *to)

{

return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);

}

static PyObject *

tee_setstate(teeobject *to, PyObject *state)

{

teedataobject *tdo;

int index;

if (!PyTuple_Check(state)) {

PyErr_SetString(PyExc_TypeError, "state is not a tuple");

return NULL;

}

if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {

return NULL;

}

if (index < 0 || index > LINKCELLS) {

PyErr_SetString(PyExc_ValueError, "Index out of range");

return NULL;

}

Py_INCREF(tdo);

Py_XSETREF(to->dataobj, tdo);

to->index = index;

Py_RETURN_NONE;

}

PyDoc_STRVAR(teeobject_doc,

"Iterator wrapped to make it copyable");

static PyMethodDef tee_methods[] = {

{"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},

{"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},

{"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},

{NULL, NULL} /* sentinel */

};

static PyTypeObject tee_type = {

PyVarObject_HEAD_INIT(NULL, 0)

"itertools._tee", /* tp_name */

sizeof(teeobject), /* tp_basicsize */

0, /* tp_itemsize */

/* methods */

(destructor)tee_dealloc, /* tp_dealloc */

0, /* tp_print */

0, /* tp_getattr */

0, /* tp_setattr */

0, /* tp_reserved */

0, /* tp_repr */

0, /* tp_as_number */

0, /* tp_as_sequence */

0, /* tp_as_mapping */

0, /* tp_hash */

0, /* tp_call */

0, /* tp_str */

0, /* tp_getattro */

0, /* tp_setattro */

0, /* tp_as_buffer */

Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */

teeobject_doc, /* tp_doc */

(traverseproc)tee_traverse, /* tp_traverse */

(inquiry)tee_clear, /* tp_clear */

0, /* tp_richcompare */

offsetof(teeobject, weakreflist), /* tp_weaklistoffset */

PyObject_SelfIter, /* tp_iter */

(iternextfunc)tee_next, /* tp_iternext */

tee_methods, /* tp_methods */

0, /* tp_members */

0, /* tp_getset */

0, /* tp_base */

0, /* tp_dict */

0, /* tp_descr_get */

0, /* tp_descr_set */

0, /* tp_dictoffset */

0, /* tp_init */

0, /* tp_alloc */

tee_new, /* tp_new */

PyObject_GC_Del, /* tp_free */

};

static PyObject *

tee(PyObject *self, PyObject *args)

{

Py_ssize_t i, n=2;

PyObject *it, *iterable, *copyable, *result;

_Py_IDENTIFIER(__copy__);

if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))

return NULL;

if (n < 0) {

PyErr_SetString(PyExc_ValueError, "n must be >= 0");

return NULL;

}

result = PyTuple_New(n);

if (result == NULL)

return NULL;

if (n == 0)

return result;

it = PyObject_GetIter(iterable);

if (it == NULL) {

Py_DECREF(result);

return NULL;

}

if (!_PyObject_HasAttrId(it, &PyId___copy__)) {

copyable = tee_fromiterable(it);

Py_DECREF(it);

if (copyable == NULL) {

Py_DECREF(result);

return NULL;

}

} else

copyable = it;

PyTuple_SET_ITEM(result, 0, copyable);

for (i=1 ; i<n ; i++) {

copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);

if (copyable == NULL) {

Py_DECREF(result);

return NULL;

}

PyTuple_SET_ITEM(result, i, copyable);

}

return result;

}

PyDoc_STRVAR(tee_doc,

"tee(iterable, n=2) --> tuple of n independent iterators.");

typedef struct {

PyObject_HEAD

PyObject *it;

PyObject *saved;

Py_ssize_t index;

int firstpass;

} cycleobject;

static PyTypeObject cycle_type;

static PyObject *

cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)

{

PyObject *it;

PyObject *iterable;

PyObject *saved;

cycleobject *lz;

if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))

return NULL;

if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))

return NULL;

/* Get iterator. */

it = PyObject_GetIter(iterable);

if (it == NULL)

return NULL;

saved = PyList_New(0);

if (saved == NULL) {

Py_DECREF(it);

return NULL;

}

/* create cycleobject structure */

lz = (cycleobject *)type->tp_alloc(type, 0);

if (lz == NULL) {

Py_DECREF(it);

Py_DECREF(saved);

return NULL;

}

lz->it = it;

lz->saved = saved;

lz->index = 0;

lz->firstpass = 0;

return (PyObject *)lz;

}

static void

cycle_dealloc(cycleobject *lz)

{

PyObject_GC_UnTrack(lz);

Py_XDECREF(lz->it);

Py_XDECREF(lz->saved);

Py_TYPE(lz)->tp_free(lz);

}

static int

cycle_traverse(cycleobject *lz, visitproc visit, void *arg)

{

if (lz->it)

Py_VISIT(lz->it);

Py_VISIT(lz->saved);

return 0;

}

static PyObject *

cycle_next(cycleobject *lz)

{

PyObject *item;

if (lz->it != NULL) {

item = PyIter_Next(lz->it);

if (item != NULL) {

if (lz->firstpass)

return item;

if (PyList_Append(lz->saved, item)) {

Py_DECREF(item);

return NULL;

}

return item;

}

/* Note: StopIteration is already cleared by PyIter_Next() */

if (PyErr_Occurred())

return NULL;

Py_CLEAR(lz->it);

}

if (PyList_GET_SIZE(lz->saved) == 0)

return NULL;

item = PyList_GET_ITEM(lz->saved, lz->index);

lz->index++;

if (lz->index >= PyList_GET_SIZE(lz->saved))

lz->index = 0;

Py_INCREF(item);

return item;

}

static PyObject *

cycle_reduce(cycleobject *lz)

{

/* Create a new cycle with the iterator tuple, then set the saved state */

if (lz->it == NULL) {

PyObject *it = PyObject_GetIter(lz->saved);

if (it == NULL)

return NULL;

if (lz->index != 0) {

_Py_IDENTIFIER(__setstate__);

PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,

"n", lz->index);

if (res == NULL) {

Py_DECREF(it);

return NULL;

}

Py_DECREF(res);

}

return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);

}

return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,

lz->firstpass);

}

static PyObject *

cycle_setstate(cycleobject *lz, PyObject *state)

{

PyObject *saved=NULL;

int firstpass;

if (!PyTuple_Check(state)) {

PyErr_SetString(PyExc_TypeError, "state is not a tuple");

return NULL;

}

if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {

return NULL;

}

Py_INCREF(saved);

Py_XSETREF(lz->saved, saved);

lz->firstpass = firstpass != 0;

lz->index = 0;

Py_RETURN_NONE;

}

static PyMethodDef cycle_methods[] = {

{"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,

reduce_doc},