Optimize PyListWriter_FromBuffer()

Author: vstinner

Include/cpython/listobject.h

@@ -83,4 +83,4 @@ PyAPI_FUNC(int) PyListWriter_FromFormat(
     ...);
 PyAPI_FUNC(int) PyListWriter_FromBuffer(
     PyListWriter *writer,
-    Py_buffer *buf);
+    const Py_buffer *buf);

Include/internal/pycore_list.h

@@ -79,6 +79,10 @@ union _PyStackRef;
 PyAPI_FUNC(PyObject *)_PyList_FromStackRefStealOnSuccess(const union _PyStackRef *src, Py_ssize_t n);
 PyAPI_FUNC(PyObject *)_PyList_AsTupleAndClear(PyListObject *v);
 
+extern void _PyListWriter_ResetSize(
+    PyListWriter *writer,
+    Py_ssize_t size);
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/test_capi/test_list.py

@@ -5,6 +5,10 @@
 from test import support
 from test.support import import_helper
 from test.support import threading_helper
+try:
+    import _testbuffer
+except ImportError:
+    _testbuffer = None
 
 _testcapi = import_helper.import_module('_testcapi')
 _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
@@ -429,6 +433,7 @@ def test_extend(self):
         # Test PyListWriter_Extend()
         self.check_extend('extend')
 
+    def test_extend_iter_error(self):
         class IterError:
             def __init__(self):
                 self.index = 1
@@ -443,7 +448,8 @@ def __next__(self):
                     raise ValueError("iter error")
                 return value
 
-        # On error, restore the writer to its previous state
+        # On error, restore the writer to its previous state:
+        # discard added items
         iterable = IterError()
         writer = self.create_writer(0)
         writer.append("before")
@@ -495,6 +501,33 @@ def test_from_buffer(self):
         writer.from_buffer(memoryview(numbers))
         self.list_equal(writer.finish(), [1.5, 3.0, 5.0])
 
+    def test_from_buffer_memoryerror(self):
+        # Inject MemoryError in the middle of the array: reset the writer
+        # to its previous state, discard added float objects
+        numbers = array.array('d', [2.0] * 128)
+        view = memoryview(numbers)
+        writer = self.create_writer(0)
+        with self.assertRaises(MemoryError):
+            with support.inject_memory_error_cm(start=len(numbers) // 2):
+                writer.from_buffer(view)
+        self.list_equal(writer.finish(), [])
+
+    @unittest.skipIf(_testbuffer is None, 'need _testbuffer')
+    def test_from_buffer_ndarray(self):
+        # PyListWriter_FromBuffer() requires 1 dimension
+        array_0d = _testbuffer.ndarray(ord(b"a"), format='B',
+                                       shape=(), strides=())
+        array_2d = _testbuffer.ndarray(list(b"abcdefgh"), format='B',
+                                       shape=(4, 2), strides=(2, 1))
+
+        for array in (array_0d, array_2d):
+            with self.subTest(array=array):
+                writer = self.create_writer(0)
+                errmsg = 'require a buffer of 1 dimension, got 2 dimentions'
+                with self.assertRaisesRegex(ValueError, errmsg):
+                    writer.from_buffer(array_2d)
+                self.list_equal(writer.finish(), [])
+
     def test_empty(self):
         list_get_allocated = _testcapi._list_get_allocated
         empty_tuple = ()

Objects/listobject.c

@@ -4553,6 +4553,28 @@ PyListWriter_Append(PyListWriter *writer, PyObject *item)
 }
 
 
+void
+_PyListWriter_ResetSize(PyListWriter *writer, Py_ssize_t size)
+{
+    assert(listwriter_check_consistency(writer));
+    PyListObject *list = _PyList_CAST(writer);
+
+    Py_ssize_t new_size = Py_SIZE(list);
+    assert(0 <= size && size <= new_size);
+
+    Py_SET_SIZE(list, size);
+
+    // Discard added items
+    PyObject **items = list->ob_item;
+    for (Py_ssize_t i = size; i < new_size; i++) {
+        Py_DECREF(items[i]);
+#ifdef Py_DEBUG
+        items[i] = NULL;
+#endif
+    }
+}
+
+
 // Similar to _list_extend() but don't lock the writer list.
 // If you update this function, update also _list_extend().
 int
@@ -4605,16 +4627,9 @@ PyListWriter_Extend(PyListWriter *writer, PyObject *iterable)
     }
 
     if (res < 0 && Py_SIZE(self) != old_size) {
-        // On error, reset the writer to its previous state
-        Py_ssize_t new_size = Py_SIZE(self);
-        Py_SET_SIZE(self, old_size);
-        PyObject **items = self->ob_item;
-        for (Py_ssize_t i = old_size; i < new_size; i++) {
-            Py_DECREF(items[i]);
-#ifdef Py_DEBUG
-            items[i] = NULL;
-#endif
-        }
+        // On error, reset the writer to its previous state:
+        // discard added items
+        _PyListWriter_ResetSize(writer, old_size);
     }
 
     return res;
@@ -4706,25 +4721,3 @@ PyListWriter_FromFormat(PyListWriter *writer, const char *format, ...)
     }
     return res;
 }
-
-
-int
-PyListWriter_FromBuffer(PyListWriter *writer, Py_buffer *view)
-{
-    extern int _PyMemoryView_ToListWriter(PyObject *self,
-                                          PyListWriter *writer);
-
-    if (view->ndim != 1) {
-        PyErr_SetString(PyExc_ValueError, "require a buffer of 1 dimension");
-        return -1;
-    }
-
-    PyObject *mview = PyMemoryView_FromBuffer(view);
-    if (mview == NULL) {
-        return -1;
-    }
-
-    int res = _PyMemoryView_ToListWriter(mview, writer);
-    Py_DECREF(mview);
-    return res;
-}

Objects/memoryobject.c

@@ -12,6 +12,7 @@
 
 #include "Python.h"
 #include "pycore_abstract.h"      // _PyIndex_Check()
+#include "pycore_list.h"          // _PyListWriter_ResetSize()
 #include "pycore_memoryobject.h"  // _PyManagedBuffer_Type
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include "pycore_strhex.h"        // _Py_strhex_with_sep()
@@ -1842,7 +1843,9 @@ unpack_single(PyMemoryViewObject *self, const char *ptr, const char *fmt)
     unsigned char uc;
     void *p;
 
-    CHECK_RELEASED_AGAIN(self);
+    if (self != NULL) {
+        CHECK_RELEASED_AGAIN(self);
+    }
 
 #if PY_LITTLE_ENDIAN
     int endian = 1;
@@ -2303,24 +2306,6 @@ tolist_base(PyMemoryViewObject *self, const char *ptr, const Py_ssize_t *shape,
     return lst;
 }
 
-static int
-tolistwriter(PyListWriter *writer, PyMemoryViewObject *self, const char *ptr,
-             const Py_ssize_t *shape, const Py_ssize_t *strides,
-             const Py_ssize_t *suboffsets, const char *fmt)
-{
-    for (Py_ssize_t i = 0; i < shape[0]; ptr+=strides[0], i++) {
-        const char *xptr = ADJUST_PTR(ptr, suboffsets, 0);
-        PyObject *item = unpack_single(self, xptr, fmt);
-        if (item == NULL) {
-            return -1;
-        }
-        if (PyListWriter_AppendSteal(writer, item) < 0) {
-            return -1;
-        }
-    }
-    return 0;
-}
-
 /* Unpack a multi-dimensional array into a nested list.
    Assumption: ndim >= 1. */
 static PyObject *
@@ -2394,21 +2379,54 @@ memoryview_tolist_impl(PyMemoryViewObject *self)
 
 
 int
-_PyMemoryView_ToListWriter(PyObject *mview, PyListWriter *writer)
+PyListWriter_FromBuffer(PyListWriter *writer, const Py_buffer *view)
 {
-    assert(PyMemoryView_Check(mview));
-    PyMemoryViewObject *self = (PyMemoryViewObject *)mview;
-    const Py_buffer *view = &self->view;
-    assert(view->ndim == 1);
-
-    const char *fmt = adjust_fmt(view);
-    if (fmt == NULL) {
+    if (view->ndim != 1) {
+        PyErr_Format(PyExc_ValueError,
+                     "require a buffer of 1 dimension, got %i dimentions",
+                     view->ndim);
         return -1;
     }
 
-    return tolistwriter(writer, self, view->buf, view->shape,
-                        view->strides, view->suboffsets,
-                        fmt);
+    const char *fmt;
+    if (view->format != NULL) {
+        fmt = adjust_fmt(view);
+        if (fmt == NULL) {
+            return -1;
+        }
+    }
+    else {
+        // Same default value than init_shared_values()
+        fmt = "B";
+    }
+
+    PyListObject *list = _PyList_CAST(writer);
+    Py_ssize_t old_writer_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;
+
+    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) {
+            goto error;
+        }
+        if (PyListWriter_AppendSteal(writer, item) < 0) {
+            goto error;
+        }
+    }
+    return 0;
+
+error:
+    if (Py_SIZE(list) != old_writer_size) {
+        // On error, reset the writer to its previous state:
+        // discard added items
+        _PyListWriter_ResetSize(writer, old_writer_size);
+    }
+    return -1;
 }