Skip to content

Commit c1e02d4

Browse files
authored
GH-98363: Presize the list for batched() (GH-98419)
1 parent 4156b2f commit c1e02d4

File tree

1 file changed

+13
-13
lines changed

1 file changed

+13
-13
lines changed

Modules/itertoolsmodule.c

+13-13
Original file line numberDiff line numberDiff line change
@@ -142,35 +142,35 @@ static PyObject *
142142
batched_next(batchedobject *bo)
143143
{
144144
Py_ssize_t i;
145+
Py_ssize_t n = bo->batch_size;
145146
PyObject *it = bo->it;
146147
PyObject *item;
147148
PyObject *result;
148149

149150
if (it == NULL) {
150151
return NULL;
151152
}
152-
result = PyList_New(0);
153+
result = PyList_New(n);
153154
if (result == NULL) {
154155
return NULL;
155156
}
156-
for (i=0 ; i < bo->batch_size ; i++) {
157+
for (i=0 ; i < n ; i++) {
157158
item = PyIter_Next(it);
158159
if (item == NULL) {
159160
break;
160161
}
161-
if (PyList_Append(result, item) < 0) {
162-
Py_DECREF(item);
163-
Py_DECREF(result);
164-
return NULL;
165-
}
166-
Py_DECREF(item);
162+
PyList_SET_ITEM(result, i, item);
167163
}
168-
if (PyList_GET_SIZE(result) > 0) {
169-
return result;
164+
if (i == 0) {
165+
Py_CLEAR(bo->it);
166+
Py_DECREF(result);
167+
return NULL;
170168
}
171-
Py_CLEAR(bo->it);
172-
Py_DECREF(result);
173-
return NULL;
169+
if (i < n) {
170+
PyObject *short_list = PyList_GetSlice(result, 0, i);
171+
Py_SETREF(result, short_list);
172+
}
173+
return result;
174174
}
175175

176176
static PyTypeObject batched_type = {

0 commit comments

Comments
 (0)