Add itertools.izip_longest().

Modules/itertoolsmodule.c

@@ -2472,6 +2472,238 @@ static PyTypeObject repeat_type = {
 	PyObject_GC_Del,		/* tp_free */
 };
 
+/* iziplongest object ************************************************************/
+
+#include "Python.h"
+
+typedef struct {
+	PyObject_HEAD
+	Py_ssize_t tuplesize;
+	Py_ssize_t numactive;	
+	PyObject *ittuple;		/* tuple of iterators */
+	PyObject *result;
+	PyObject *fillvalue;
+	PyObject *filler;		/* repeat(fillvalue) */
+} iziplongestobject;
+
+static PyTypeObject iziplongest_type;
+
+static PyObject *
+izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+	iziplongestobject *lz;
+	Py_ssize_t i;
+	PyObject *ittuple;  /* tuple of iterators */
+	PyObject *result;
+	PyObject *fillvalue = Py_None;
+	PyObject *filler;
+	Py_ssize_t tuplesize = PySequence_Length(args);
+
+        if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
+                fillvalue = PyDict_GetItemString(kwds, "fillvalue");
+                if (fillvalue == NULL  ||  PyDict_Size(kwds) > 1) {
+                        PyErr_SetString(PyExc_TypeError,
+				"izip_longest() got an unexpected keyword argument");
+                        return NULL;                      
+                }
+        }
+
+	/* args must be a tuple */
+	assert(PyTuple_Check(args));
+
+	/* obtain iterators */
+	ittuple = PyTuple_New(tuplesize);
+	if (ittuple == NULL)
+		return NULL;
+	for (i=0; i < tuplesize; ++i) {
+		PyObject *item = PyTuple_GET_ITEM(args, i);
+		PyObject *it = PyObject_GetIter(item);
+		if (it == NULL) {
+			if (PyErr_ExceptionMatches(PyExc_TypeError))
+				PyErr_Format(PyExc_TypeError,
+				    "izip_longest argument #%zd must support iteration",
+				    i+1);
+			Py_DECREF(ittuple);
+			return NULL;
+		}
+		PyTuple_SET_ITEM(ittuple, i, it);
+	}
+
+	filler = PyObject_CallFunctionObjArgs((PyObject *)(&repeat_type), fillvalue, NULL);
+	if (filler == NULL) {
+		Py_DECREF(ittuple);
+		return NULL;
+	}
+
+	/* create a result holder */
+	result = PyTuple_New(tuplesize);
+	if (result == NULL) {
+		Py_DECREF(ittuple);
+		Py_DECREF(filler);
+		return NULL;
+	}
+	for (i=0 ; i < tuplesize ; i++) {
+		Py_INCREF(Py_None);
+		PyTuple_SET_ITEM(result, i, Py_None);
+	}
+
+	/* create iziplongestobject structure */
+	lz = (iziplongestobject *)type->tp_alloc(type, 0);
+	if (lz == NULL) {
+		Py_DECREF(ittuple);
+		Py_DECREF(filler);
+		Py_DECREF(result);
+		return NULL;
+	}
+	lz->ittuple = ittuple;
+	lz->tuplesize = tuplesize;
+	lz->numactive = tuplesize;
+	lz->result = result;
+	Py_INCREF(fillvalue);
+	lz->fillvalue = fillvalue;
+	Py_INCREF(filler);
+	lz->filler = filler;			/* XXX */
+	return (PyObject *)lz;
+}
+
+static void
+izip_longest_dealloc(iziplongestobject *lz)
+{
+	PyObject_GC_UnTrack(lz);
+	Py_XDECREF(lz->ittuple);
+	Py_XDECREF(lz->result);
+	Py_XDECREF(lz->fillvalue);
+	Py_XDECREF(lz->filler);
+	lz->ob_type->tp_free(lz);
+}
+
+static int
+izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
+{
+	Py_VISIT(lz->ittuple);
+	Py_VISIT(lz->result);
+	Py_VISIT(lz->fillvalue);
+	Py_VISIT(lz->filler);
+	return 0;
+}
+
+static PyObject *
+izip_longest_next(iziplongestobject *lz)
+{
+	Py_ssize_t i;
+	Py_ssize_t tuplesize = lz->tuplesize;
+	PyObject *result = lz->result;
+	PyObject *it;
+	PyObject *item;
+	PyObject *olditem;
+
+	if (tuplesize == 0)
+		return NULL;
+	if (result->ob_refcnt == 1) {
+		Py_INCREF(result);
+		for (i=0 ; i < tuplesize ; i++) {
+			it = PyTuple_GET_ITEM(lz->ittuple, i);
+			assert(PyIter_Check(it));
+			item = (*it->ob_type->tp_iternext)(it);
+			if (item == NULL) {
+				if (lz->numactive <= 1) {
+					Py_DECREF(result);
+					return NULL;
+				} else {
+					Py_INCREF(lz->filler);
+					PyTuple_SET_ITEM(lz->ittuple, i, lz->filler);
+					Py_INCREF(lz->fillvalue);
+					item = lz->fillvalue;
+					Py_DECREF(it);
+					lz->numactive -= 1;
+				}
+			}
+			olditem = PyTuple_GET_ITEM(result, i);
+			PyTuple_SET_ITEM(result, i, item);
+			Py_DECREF(olditem);
+		}
+	} else {
+		result = PyTuple_New(tuplesize);
+		if (result == NULL)
+			return NULL;
+		for (i=0 ; i < tuplesize ; i++) {
+			it = PyTuple_GET_ITEM(lz->ittuple, i);
+			assert(PyIter_Check(it));
+			item = (*it->ob_type->tp_iternext)(it);
+			if (item == NULL) {
+				if (lz->numactive <= 1) {
+					Py_DECREF(result);
+					return NULL;
+				} else {
+					Py_INCREF(lz->filler);
+					PyTuple_SET_ITEM(lz->ittuple, i, lz->filler);
+					Py_INCREF(lz->fillvalue);
+					item = lz->fillvalue;
+					Py_DECREF(it);
+					lz->numactive -= 1;
+				}
+			}
+			PyTuple_SET_ITEM(result, i, item);
+		}
+	}
+	return result;
+}
+
+PyDoc_STRVAR(izip_longest_doc,
+"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
+\n\
+Return an izip_longest object whose .next() method returns a tuple where\n\
+the i-th element comes from the i-th iterable argument.  The .next()\n\
+method continues until the longest iterable in the argument sequence\n\
+is exhausted and then it raises StopIteration.  When the shorter iterables\n\
+are exhausted, the fillvalue is substituted in their place.  The fillvalue\n\
+defaults to None or can be specified by a keyword argument.\n\
+");
+
+static PyTypeObject iziplongest_type = {
+	PyObject_HEAD_INIT(NULL)
+	0,				/* ob_size */
+	"itertools.izip_longest",	/* tp_name */
+	sizeof(iziplongestobject),	/* tp_basicsize */
+	0,				/* tp_itemsize */
+	/* methods */
+	(destructor)izip_longest_dealloc,	/* tp_dealloc */
+	0,				/* tp_print */
+	0,				/* tp_getattr */
+	0,				/* tp_setattr */
+	0,				/* tp_compare */
+	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 */
+	izip_longest_doc,			/* tp_doc */
+	(traverseproc)izip_longest_traverse,    /* tp_traverse */
+	0,				/* tp_clear */
+	0,				/* tp_richcompare */
+	0,				/* tp_weaklistoffset */
+	PyObject_SelfIter,		/* tp_iter */
+	(iternextfunc)izip_longest_next,	/* tp_iternext */
+	0,				/* 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 */
+	izip_longest_new,			/* tp_new */
+	PyObject_GC_Del,		/* tp_free */
+};
 
 /* module level code ********************************************************/
 
@@ -2485,6 +2717,7 @@ repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
 \n\
 Iterators terminating on the shortest input sequence:\n\
 izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
+izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
 ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
 ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
 islice(seq, [start,] stop [, step]) --> elements from\n\
@@ -2522,6 +2755,7 @@ inititertools(void)
 		&ifilterfalse_type,
 		&count_type,
 		&izip_type,
+		&iziplongest_type,                
 		&repeat_type,
 		&groupby_type,
 		NULL