diff options
author | Raymond Hettinger <rhettinger@users.noreply.github.com> | 2022-10-20 07:28:17 (GMT) |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-10-20 07:28:17 (GMT) |
commit | c1e02d4e4e09dfc9d57c4d36eca1bbf312d4162d (patch) | |
tree | 84f634f9d63a717de274179e4535708a3cfc74f0 | |
parent | 4156b2fc1378531bdd9459cd131fb9fa25a5af34 (diff) | |
download | cpython-c1e02d4e4e09dfc9d57c4d36eca1bbf312d4162d.zip cpython-c1e02d4e4e09dfc9d57c4d36eca1bbf312d4162d.tar.gz cpython-c1e02d4e4e09dfc9d57c4d36eca1bbf312d4162d.tar.bz2 |
GH-98363: Presize the list for batched() (GH-98419)
-rw-r--r-- | Modules/itertoolsmodule.c | 26 |
1 files changed, 13 insertions, 13 deletions
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index a5bbba1..868e8a8 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -142,6 +142,7 @@ static PyObject * batched_next(batchedobject *bo) { Py_ssize_t i; + Py_ssize_t n = bo->batch_size; PyObject *it = bo->it; PyObject *item; PyObject *result; @@ -149,28 +150,27 @@ batched_next(batchedobject *bo) if (it == NULL) { return NULL; } - result = PyList_New(0); + result = PyList_New(n); if (result == NULL) { return NULL; } - for (i=0 ; i < bo->batch_size ; i++) { + for (i=0 ; i < n ; i++) { item = PyIter_Next(it); if (item == NULL) { break; } - if (PyList_Append(result, item) < 0) { - Py_DECREF(item); - Py_DECREF(result); - return NULL; - } - Py_DECREF(item); + PyList_SET_ITEM(result, i, item); } - if (PyList_GET_SIZE(result) > 0) { - return result; + if (i == 0) { + Py_CLEAR(bo->it); + Py_DECREF(result); + return NULL; } - Py_CLEAR(bo->it); - Py_DECREF(result); - return NULL; + if (i < n) { + PyObject *short_list = PyList_GetSlice(result, 0, i); + Py_SETREF(result, short_list); + } + return result; } static PyTypeObject batched_type = { |