Make built-in zip() equal to itertools.izip().
I mea, *really* equal -- for now, the implementation just imports
itertools. :-)
The only other changes necessary were various unit tests that were
assuming zip() returns a real list. No "real" code made this assumption.
diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 94420f8..6ca2a28 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -1855,116 +1855,32 @@
static PyObject*
builtin_zip(PyObject *self, PyObject *args)
{
- PyObject *ret;
- const Py_ssize_t itemsize = PySequence_Length(args);
- Py_ssize_t i;
- PyObject *itlist; /* tuple of iterators */
- Py_ssize_t len; /* guess at result length */
+ PyObject *itertools = NULL, *izip = NULL, *result = NULL;
- if (itemsize == 0)
- return PyList_New(0);
-
- /* args must be a tuple */
- assert(PyTuple_Check(args));
-
- /* Guess at result length: the shortest of the input lengths.
- If some argument refuses to say, we refuse to guess too, lest
- an argument like xrange(sys.maxint) lead us astray.*/
- len = -1; /* unknown */
- for (i = 0; i < itemsize; ++i) {
- PyObject *item = PyTuple_GET_ITEM(args, i);
- Py_ssize_t thislen = _PyObject_LengthHint(item);
- if (thislen < 0) {
- if (!PyErr_ExceptionMatches(PyExc_TypeError) &&
- !PyErr_ExceptionMatches(PyExc_AttributeError)) {
- return NULL;
- }
- PyErr_Clear();
- len = -1;
- break;
- }
- else if (len < 0 || thislen < len)
- len = thislen;
- }
-
- /* allocate result list */
- if (len < 0)
- len = 10; /* arbitrary */
- if ((ret = PyList_New(len)) == NULL)
+ itertools = PyImport_ImportModule("itertools");
+ if (itertools == NULL)
return NULL;
+
+ izip = PyObject_GetAttrString(itertools, "izip");
+ if (izip == NULL)
+ goto done;
- /* obtain iterators */
- itlist = PyTuple_New(itemsize);
- if (itlist == NULL)
- goto Fail_ret;
- for (i = 0; i < itemsize; ++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,
- "zip argument #%zd must support iteration",
- i+1);
- goto Fail_ret_itlist;
- }
- PyTuple_SET_ITEM(itlist, i, it);
- }
+ result = PyObject_Call(izip, args, NULL);
- /* build result into ret list */
- for (i = 0; ; ++i) {
- int j;
- PyObject *next = PyTuple_New(itemsize);
- if (!next)
- goto Fail_ret_itlist;
-
- for (j = 0; j < itemsize; j++) {
- PyObject *it = PyTuple_GET_ITEM(itlist, j);
- PyObject *item = PyIter_Next(it);
- if (!item) {
- if (PyErr_Occurred()) {
- Py_DECREF(ret);
- ret = NULL;
- }
- Py_DECREF(next);
- Py_DECREF(itlist);
- goto Done;
- }
- PyTuple_SET_ITEM(next, j, item);
- }
-
- if (i < len)
- PyList_SET_ITEM(ret, i, next);
- else {
- int status = PyList_Append(ret, next);
- Py_DECREF(next);
- ++len;
- if (status < 0)
- goto Fail_ret_itlist;
- }
- }
-
-Done:
- if (ret != NULL && i < len) {
- /* The list is too big. */
- if (PyList_SetSlice(ret, i, len, NULL) < 0)
- return NULL;
- }
- return ret;
-
-Fail_ret_itlist:
- Py_DECREF(itlist);
-Fail_ret:
- Py_DECREF(ret);
- return NULL;
+ done:
+ Py_XDECREF(itertools);
+ Py_XDECREF(izip);
+ return result;
}
PyDoc_STRVAR(zip_doc,
-"zip(seq1 [, seq2 [...]]) -> [(seq1[0], seq2[0] ...), (...)]\n\
+"zip(it1 [, it2 [...]]) -> iter([(it1[0], it2[0] ...), ...])\n\
\n\
-Return a list of tuples, where each tuple contains the i-th element\n\
-from each of the argument sequences. The returned list is truncated\n\
-in length to the length of the shortest argument sequence.");
+Return an iterator yielding tuples, where each tuple contains the\n\
+corresponding element from each of the argument iterables.\n\
+The returned iterator ends when the shortest argument iterable is exhausted.\n\
+NOTE: This is implemented using itertools.izip().");
static PyMethodDef builtin_methods[] = {