Resize the list once in PyListWriter_FromBuffer()

Add _PyListWriter_SET_ITEM_STEAL() static inline function.

Author: vstinner

Doc/c-api/list.rst

@@ -304,7 +304,7 @@ On Free Threading, it's more efficient than the ``PyList`` C API to create a
    On error, set an exception, leave the writer unchanged, and return ``-1``.
 
    If consuming *iterable* raises an exception, the writer is restored to its
-   previous state: added items are discarded.
+   previous state: discard added items.
 
 .. c:function:: int PyListWriter_ExtendArray(PyListWriter *writer, PyObject *const *items, Py_ssize_t nitem)
 
@@ -330,6 +330,8 @@ On Free Threading, it's more efficient than the ``PyList`` C API to create a
    On success, return ``0``.
    On error, set an exception, leave the writer unchanged, and return ``-1``.
 
+   On error, the writer is restored to its previous state: discard added items.
+
 .. c:function:: int PyListWriter_FromBuffer(PyListWriter *writer, Py_buffer *view)
 
    Add values from a buffer.
@@ -338,3 +340,5 @@ On Free Threading, it's more efficient than the ``PyList`` C API to create a
 
    On success, return ``0``.
    On error, set an exception, leave the writer unchanged, and return ``-1``.
+
+   On error, the writer is restored to its previous state: discard added items.

Include/internal/pycore_list.h

@@ -79,10 +79,37 @@ union _PyStackRef;
 PyAPI_FUNC(PyObject *)_PyList_FromStackRefStealOnSuccess(const union _PyStackRef *src, Py_ssize_t n);
 PyAPI_FUNC(PyObject *)_PyList_AsTupleAndClear(PyListObject *v);
 
+#ifdef Py_DEBUG
+#  define _PyListWriter_INIT_ITEMS 1
+#else
+#  define _PyListWriter_INIT_ITEMS 0
+#endif
+
+extern int _PyListWriter_Resize(
+    PyListWriter *writer,
+    Py_ssize_t new_size);
 extern void _PyListWriter_ResetSize(
     PyListWriter *writer,
     Py_ssize_t size);
 
+
+// The caller must resize the list so it's big enough to set this item.
+// A writer list item must be set exactly once.
+static inline void
+_PyListWriter_SET_ITEM_STEAL(PyListWriter *writer, Py_ssize_t index,
+                             PyObject *item)
+{
+    assert(item != NULL);
+
+    PyListObject *list = (PyListObject*)writer;
+    assert(0 <= index && index < Py_SIZE(list));
+#if _PyListWriter_INIT_ITEMS
+    assert(list->ob_item[index] == NULL);
+#endif
+
+    list->ob_item[index] = item;
+}
+
 #ifdef __cplusplus
 }
 #endif

Objects/listobject.c

@@ -4449,7 +4449,7 @@ PyListWriter_Create(Py_ssize_t size)
     // since the list object is not tracked by the GC, and the writer caller
     // must initialize all items.
 
-#ifdef Py_DEBUG
+#if _PyListWriter_INIT_ITEMS
     // In debug mode, set all items to NULL to be able to check that they are
     // non-NULL in PyListWriter_Finish().
     if (1 <= size) {
@@ -4498,7 +4498,7 @@ PyListWriter_Finish(PyListWriter *writer)
         // it's left unchanged on memory allocation failure.
     }
 
-#ifdef Py_DEBUG
+#if _PyListWriter_INIT_ITEMS
     // Make sure that all items have been initialized
     Py_ssize_t size = Py_SIZE(list);
     PyObject **list_items = list->ob_item;
@@ -4513,13 +4513,21 @@ PyListWriter_Finish(PyListWriter *writer)
 }
 
 
-static int
-listwriter_resize(PyListWriter *writer, Py_ssize_t new_size)
+int
+_PyListWriter_Resize(PyListWriter *writer, Py_ssize_t new_size)
 {
     PyListObject *list = _PyList_CAST(writer);
+    Py_ssize_t old_size = Py_SIZE(list);
 
     // On Free Threading, it can call _PyObject_GC_SET_SHARED() on the list
-    return list_resize(list, new_size);
+    int res = list_resize(list, new_size);
+#if _PyListWriter_INIT_ITEMS
+    if (res == 0 && new_size > old_size) {
+        memset(&list->ob_item[old_size], 0,
+               (new_size - old_size) * sizeof(PyObject*));
+    }
+#endif
+    return res;
 }
 
 
@@ -4529,19 +4537,19 @@ PyListWriter_AppendSteal(PyListWriter *writer, PyObject *item)
     assert(listwriter_check_consistency(writer));
     PyListObject *list = _PyList_CAST(writer);
 
-    Py_ssize_t len = Py_SIZE(list);
+    Py_ssize_t index = Py_SIZE(list);
     Py_ssize_t allocated = list->allocated;
-    if (len < allocated) {
-        list->ob_item[len] = item;
-        Py_SET_SIZE(list, len + 1);
+    if (index < allocated) {
+        Py_SET_SIZE(list, index + 1);
+        list->ob_item[index] = item;
         return 0;
     }
 
-    if (listwriter_resize(writer, len + 1) < 0) {
+    if (_PyListWriter_Resize(writer, index + 1) < 0) {
         Py_DECREF(item);
         return -1;
     }
-    list->ob_item[len] = item;
+    _PyListWriter_SET_ITEM_STEAL(writer, index, item);
     return 0;
 }
 
@@ -4568,7 +4576,7 @@ _PyListWriter_ResetSize(PyListWriter *writer, Py_ssize_t size)
     PyObject **items = list->ob_item;
     for (Py_ssize_t i = size; i < new_size; i++) {
         Py_DECREF(items[i]);
-#ifdef Py_DEBUG
+#if _PyListWriter_INIT_ITEMS
         items[i] = NULL;
 #endif
     }
@@ -4652,10 +4660,9 @@ listwriter_extend_array(PyListWriter *writer,
 
     PyListObject *list = _PyList_CAST(writer);
     Py_ssize_t size = Py_SIZE(list);
-    Py_ssize_t allocated = list->allocated;
     Py_ssize_t new_size = size + nitem;
-    if (new_size > allocated) {
-        if (listwriter_resize(writer, new_size) < 0) {
+    if (new_size > list->allocated) {
+        if (_PyListWriter_Resize(writer, new_size) < 0) {
             if (steal) {
                 for (Py_ssize_t i = 0; i < nitem; i++) {
                     Py_DECREF(items[i]);

Objects/memoryobject.c

@@ -2401,30 +2401,41 @@ PyListWriter_FromBuffer(PyListWriter *writer, const Py_buffer *view)
     }
 
     PyListObject *list = _PyList_CAST(writer);
-    Py_ssize_t old_writer_size = Py_SIZE(list);
+    Py_ssize_t old_size = Py_SIZE(list);
 
     const char *ptr = view->buf;
     const Py_ssize_t shape = view->shape ? view->shape[0] : view->len;
     const Py_ssize_t strides = view->strides ? view->strides[0] : view->itemsize;
     const Py_ssize_t *suboffsets = view->suboffsets;
 
+    Py_ssize_t new_size = old_size + shape;
+    if (new_size > list->allocated) {
+        if (_PyListWriter_Resize(writer, new_size) < 0) {
+            return -1;
+        }
+    }
+    else {
+        Py_SET_SIZE(list, new_size);
+    }
+
     for (Py_ssize_t i = 0; i < shape; ptr += strides, i++) {
         const char *xptr = ADJUST_PTR(ptr, suboffsets, 0);
+
         PyObject *item = unpack_single(NULL, xptr, fmt);
         if (item == NULL) {
+            Py_SET_SIZE(list, old_size + i);
             goto error;
         }
-        if (PyListWriter_AppendSteal(writer, item) < 0) {
-            goto error;
-        }
+
+        _PyListWriter_SET_ITEM_STEAL(writer, old_size + i, item);
     }
     return 0;
 
 error:
-    if (Py_SIZE(list) != old_writer_size) {
+    if (Py_SIZE(list) != old_size) {
         // On error, reset the writer to its previous state:
         // discard added items
-        _PyListWriter_ResetSize(writer, old_writer_size);
+        _PyListWriter_ResetSize(writer, old_size);
     }
     return -1;
 }