builtin_zip(): Take a good guess at how big the result list will be,

and allocate it in one gulp.

This isn't a bugfix, it's just a minor optimization that may or may not
pay off.

diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
index 13ed5bd4a8b549bb4e0564cb156434caa983a18b..c9f696166c4644377ac0b6446e65e3f4b7f99705 100644 (file)
--- a/Lib/test/test_iter.py
+++ b/Lib/test/test_iter.py
@@ -467,6 +467,34 @@ class TestCase(unittest.TestCase):
             except OSError:
                 pass
 
+        self.assertEqual(zip(xrange(5)), [(i,) for i in range(5)])
+
+        # Classes that lie about their lengths.
+        class NoGuessLen5:
+            def __getitem__(self, i):
+                if i >= 5:
+                    raise IndexError
+                return i
+
+        class Guess3Len5(NoGuessLen5):
+            def __len__(self):
+                return 3
+
+        class Guess30Len5(NoGuessLen5):
+            def __len__(self):
+                return 30
+
+        self.assertEqual(len(Guess3Len5()), 3)
+        self.assertEqual(len(Guess30Len5()), 30)
+        self.assertEqual(zip(NoGuessLen5()), zip(range(5)))
+        self.assertEqual(zip(Guess3Len5()), zip(range(5)))
+        self.assertEqual(zip(Guess30Len5()), zip(range(5)))
+
+        expected = [(i, i) for i in range(5)]
+        for x in NoGuessLen5(), Guess3Len5(), Guess30Len5():
+            for y in NoGuessLen5(), Guess3Len5(), Guess30Len5():
+                self.assertEqual(zip(x, y), expected)
+
     # Test reduces()'s use of iterators.
     def test_builtin_reduce(self):
         from operator import add

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index aa65ebca410d628213d29e1956e8d1d073711aeb..fb6810edb8d94de2e2acb20ca335f3486214ab8b 100644 (file)
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -381,7 +381,7 @@ builtin_compile(PyObject *self, PyObject *args)
        int supplied_flags = 0;
        PyCompilerFlags cf;
 
-       if (!PyArg_ParseTuple(args, "sss|ii:compile", &str, &filename, 
+       if (!PyArg_ParseTuple(args, "sss|ii:compile", &str, &filename,
                              &startstr, &supplied_flags, &dont_inherit))
                return NULL;
 
@@ -1128,7 +1128,7 @@ min_max(PyObject *args, int op)
                v = args;
        else if (!PyArg_ParseTuple(args, "O:min/max", &v))
                return NULL;
-       
+
        it = PyObject_GetIter(v);
        if (it == NULL)
                return NULL;
@@ -1704,9 +1704,10 @@ static PyObject*
 builtin_zip(PyObject *self, PyObject *args)
 {
        PyObject *ret;
-       int itemsize = PySequence_Length(args);
+       const int itemsize = PySequence_Length(args);
        int i;
        PyObject *itlist;  /* tuple of iterators */
+       int len;           /* guess at result length */
 
        if (itemsize < 1) {
                PyErr_SetString(PyExc_TypeError,
@@ -1716,8 +1717,21 @@ builtin_zip(PyObject *self, PyObject *args)
        /* args must be a tuple */
        assert(PyTuple_Check(args));
 
+       /* Guess at result length:  the shortest of the input lengths. */
+       len = -1;       /* unknown */
+       for (i = 0; i < itemsize; ++i) {
+               PyObject *item = PyTuple_GET_ITEM(args, i);
+               int thislen = PySequence_Length(item);
+               if (thislen < 0)
+                       PyErr_Clear();
+               else if (len < 0 || thislen < len)
+                       len = thislen;
+       }
+
        /* allocate result list */
-       if ((ret = PyList_New(0)) == NULL)
+       if (len < 0)
+               len = 10;       /* arbitrary */
+       if ((ret = PyList_New(len)) == NULL)
                return NULL;
 
        /* obtain iterators */
@@ -1738,14 +1752,14 @@ builtin_zip(PyObject *self, PyObject *args)
        }
 
        /* build result into ret list */
-       for (;;) {
-               int status;
+       for (i = 0; ; ++i) {
+               int j;
                PyObject *next = PyTuple_New(itemsize);
                if (!next)
                        goto Fail_ret_itlist;
 
-               for (i = 0; i < itemsize; i++) {
-                       PyObject *it = PyTuple_GET_ITEM(itlist, i);
+               for (j = 0; j < itemsize; j++) {
+                       PyObject *it = PyTuple_GET_ITEM(itlist, j);
                        PyObject *item = PyIter_Next(it);
                        if (!item) {
                                if (PyErr_Occurred()) {
@@ -1754,16 +1768,29 @@ builtin_zip(PyObject *self, PyObject *args)
                                }
                                Py_DECREF(next);
                                Py_DECREF(itlist);
-                               return ret;
+                               goto Done;
                        }
-                       PyTuple_SET_ITEM(next, i, item);
+                       PyTuple_SET_ITEM(next, j, item);
                }
 
-               status = PyList_Append(ret, next);
-               Py_DECREF(next);
-               if (status < 0)
-                       goto Fail_ret_itlist;
+               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);
@@ -1864,7 +1891,7 @@ _PyBuiltin_Init(void)
        SETBUILTIN("complex",           &PyComplex_Type);
 #endif
        SETBUILTIN("dict",              &PyDict_Type);
-       SETBUILTIN("enumerate",         &PyEnum_Type);  
+       SETBUILTIN("enumerate",         &PyEnum_Type);
        SETBUILTIN("float",             &PyFloat_Type);
        SETBUILTIN("property",          &PyProperty_Type);
        SETBUILTIN("int",               &PyInt_Type);