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.
This commit is contained in:
Tim Peters 2002-04-29 21:27:32 +00:00
parent 541703b18f
commit 67d687a114
2 changed files with 70 additions and 15 deletions

View file

@ -467,6 +467,34 @@ class TestCase(unittest.TestCase):
except OSError: except OSError:
pass 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. # Test reduces()'s use of iterators.
def test_builtin_reduce(self): def test_builtin_reduce(self):
from operator import add from operator import add

View file

@ -1704,9 +1704,10 @@ static PyObject*
builtin_zip(PyObject *self, PyObject *args) builtin_zip(PyObject *self, PyObject *args)
{ {
PyObject *ret; PyObject *ret;
int itemsize = PySequence_Length(args); const int itemsize = PySequence_Length(args);
int i; int i;
PyObject *itlist; /* tuple of iterators */ PyObject *itlist; /* tuple of iterators */
int len; /* guess at result length */
if (itemsize < 1) { if (itemsize < 1) {
PyErr_SetString(PyExc_TypeError, PyErr_SetString(PyExc_TypeError,
@ -1716,8 +1717,21 @@ builtin_zip(PyObject *self, PyObject *args)
/* args must be a tuple */ /* args must be a tuple */
assert(PyTuple_Check(args)); 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 */ /* allocate result list */
if ((ret = PyList_New(0)) == NULL) if (len < 0)
len = 10; /* arbitrary */
if ((ret = PyList_New(len)) == NULL)
return NULL; return NULL;
/* obtain iterators */ /* obtain iterators */
@ -1738,14 +1752,14 @@ builtin_zip(PyObject *self, PyObject *args)
} }
/* build result into ret list */ /* build result into ret list */
for (;;) { for (i = 0; ; ++i) {
int status; int j;
PyObject *next = PyTuple_New(itemsize); PyObject *next = PyTuple_New(itemsize);
if (!next) if (!next)
goto Fail_ret_itlist; goto Fail_ret_itlist;
for (i = 0; i < itemsize; i++) { for (j = 0; j < itemsize; j++) {
PyObject *it = PyTuple_GET_ITEM(itlist, i); PyObject *it = PyTuple_GET_ITEM(itlist, j);
PyObject *item = PyIter_Next(it); PyObject *item = PyIter_Next(it);
if (!item) { if (!item) {
if (PyErr_Occurred()) { if (PyErr_Occurred()) {
@ -1754,16 +1768,29 @@ builtin_zip(PyObject *self, PyObject *args)
} }
Py_DECREF(next); Py_DECREF(next);
Py_DECREF(itlist); Py_DECREF(itlist);
return ret; goto Done;
} }
PyTuple_SET_ITEM(next, i, item); PyTuple_SET_ITEM(next, j, item);
} }
status = PyList_Append(ret, next); if (i < len)
PyList_SET_ITEM(ret, i, next);
else {
int status = PyList_Append(ret, next);
Py_DECREF(next); Py_DECREF(next);
++len;
if (status < 0) if (status < 0)
goto Fail_ret_itlist; 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: Fail_ret_itlist:
Py_DECREF(itlist); Py_DECREF(itlist);