Skip to content

Commit 0e05aa4

Browse files
committed
Cleaner handling of empty result
1 parent 88fc1b1 commit 0e05aa4

File tree

1 file changed

+10
-9
lines changed

1 file changed

+10
-9
lines changed

Modules/itertoolsmodule.c

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -142,37 +142,38 @@ 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(bo->batch_size);
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
}
161162
PyList_SET_ITEM(result, i, item);
162163
}
163-
if (i > 0 && i < bo->batch_size) {
164+
if (i == 0) {
165+
Py_CLEAR(bo->it);
166+
Py_DECREF(result);
167+
return NULL;
168+
}
169+
if (i < bo->batch_size) {
164170
PyObject *short_list = PyList_GetSlice(result, 0, i);
165171
Py_SETREF(result, short_list);
166172
if (result == NULL) {
167173
return NULL;
168174
}
169175
}
170-
if (i > 0) {
171-
return result;
172-
}
173-
Py_CLEAR(bo->it);
174-
Py_DECREF(result);
175-
return NULL;
176+
return result;
176177
}
177178

178179
static PyTypeObject batched_type = {

0 commit comments

Comments
 (0)