From ea9511e781dea7b25d8929332534afefdb6dadcd Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Tue, 18 Oct 2022 18:01:56 -0500 Subject: [PATCH 1/6] Presize the list for batched() --- Modules/itertoolsmodule.c | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index a5bbba14c2808e..1ce35335fc0a45 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -149,7 +149,7 @@ batched_next(batchedobject *bo) if (it == NULL) { return NULL; } - result = PyList_New(0); + result = PyList_New(bo->batch_size); if (result == NULL) { return NULL; } @@ -158,12 +158,14 @@ batched_next(batchedobject *bo) if (item == NULL) { break; } - if (PyList_Append(result, item) < 0) { - Py_DECREF(item); - Py_DECREF(result); + PyList_SET_ITEM(result, i, item); + } + if (i < bo->batch_size) { + PyObject *short_list = PyList_GetSlice(result, 0, i); + Py_SETREF(result, short_list); + if (result == NULL) { return NULL; } - Py_DECREF(item); } if (PyList_GET_SIZE(result) > 0) { return result; From 88fc1b105a465370c1eee0e6e2a431a612e33059 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Tue, 18 Oct 2022 19:29:30 -0500 Subject: [PATCH 2/6] Fast path for the common case (zero new items). --- Modules/itertoolsmodule.c | 4 ++-- Objects/listobject.c | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 1ce35335fc0a45..caf222cf6a2ef3 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -160,14 +160,14 @@ batched_next(batchedobject *bo) } PyList_SET_ITEM(result, i, item); } - if (i < bo->batch_size) { + if (i > 0 && i < bo->batch_size) { PyObject *short_list = PyList_GetSlice(result, 0, i); Py_SETREF(result, short_list); if (result == NULL) { return NULL; } } - if (PyList_GET_SIZE(result) > 0) { + if (i > 0) { return result; } Py_CLEAR(bo->it); diff --git a/Objects/listobject.c b/Objects/listobject.c index 2cb3de9f147cc7..8516871b63e650 100644 --- a/Objects/listobject.c +++ b/Objects/listobject.c @@ -7,6 +7,7 @@ #include "pycore_object.h" // _PyObject_GC_TRACK() #include "pycore_tuple.h" // _PyTuple_FromArray() #include +#include /*[clinic input] class list "PyListObject *" "&PyList_Type" From 0e05aa486104ca6639e43f7e93fca7708696dacc Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Tue, 18 Oct 2022 19:42:31 -0500 Subject: [PATCH 3/6] Cleaner handling of empty result --- Modules/itertoolsmodule.c | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index caf222cf6a2ef3..4a5bd3b86b4691 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,30 +150,30 @@ batched_next(batchedobject *bo) if (it == NULL) { return NULL; } - result = PyList_New(bo->batch_size); + 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; } PyList_SET_ITEM(result, i, item); } - if (i > 0 && i < bo->batch_size) { + if (i == 0) { + Py_CLEAR(bo->it); + Py_DECREF(result); + return NULL; + } + if (i < bo->batch_size) { PyObject *short_list = PyList_GetSlice(result, 0, i); Py_SETREF(result, short_list); if (result == NULL) { return NULL; } } - if (i > 0) { - return result; - } - Py_CLEAR(bo->it); - Py_DECREF(result); - return NULL; + return result; } static PyTypeObject batched_type = { From 646630efe3f69ef626d263162de684f1b0e6f7a2 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Tue, 18 Oct 2022 19:43:22 -0500 Subject: [PATCH 4/6] Let NULL error flow through the return --- Modules/itertoolsmodule.c | 3 --- 1 file changed, 3 deletions(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 4a5bd3b86b4691..d2798d007fd8c8 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -169,9 +169,6 @@ batched_next(batchedobject *bo) if (i < bo->batch_size) { PyObject *short_list = PyList_GetSlice(result, 0, i); Py_SETREF(result, short_list); - if (result == NULL) { - return NULL; - } } return result; } From 5e38c096ff98b319c4f45a3d857f8187f0aa98a2 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Tue, 18 Oct 2022 19:44:41 -0500 Subject: [PATCH 5/6] Use n variable consistently throughout --- Modules/itertoolsmodule.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index d2798d007fd8c8..868e8a8b384f1b 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -166,7 +166,7 @@ batched_next(batchedobject *bo) Py_DECREF(result); return NULL; } - if (i < bo->batch_size) { + if (i < n) { PyObject *short_list = PyList_GetSlice(result, 0, i); Py_SETREF(result, short_list); } From fd5b50871810e8228d6596d7681beb883d9ecebd Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Tue, 18 Oct 2022 19:47:33 -0500 Subject: [PATCH 6/6] Remove spurious include directive --- Objects/listobject.c | 1 - 1 file changed, 1 deletion(-) diff --git a/Objects/listobject.c b/Objects/listobject.c index 8516871b63e650..2cb3de9f147cc7 100644 --- a/Objects/listobject.c +++ b/Objects/listobject.c @@ -7,7 +7,6 @@ #include "pycore_object.h" // _PyObject_GC_TRACK() #include "pycore_tuple.h" // _PyTuple_FromArray() #include -#include /*[clinic input] class list "PyListObject *" "&PyList_Type"