Finish() shrinks the list

Add test_empty() and test_finish_shrink() tests. Add tests on
MemoryError.

Move LIST_SMALL_ALLOCATED to pycore_list.h as
_Py_LIST_SMALL_ALLOCATED.

Author: vstinner

Include/internal/pycore_list.h

@@ -12,6 +12,8 @@ extern "C" {
 #include "pycore_stackref.h"
 #endif
 
+#define _Py_LIST_SMALL_ALLOCATED 32
+
 PyAPI_FUNC(PyObject*) _PyList_Extend(PyListObject *, PyObject *);
 PyAPI_FUNC(PyObject) *_PyList_SliceSubscript(PyObject*, PyObject*);
 PyAPI_FUNC(PyObject *) _PyList_BinarySlice(PyObject *, PyObject *, PyObject *);

Lib/test/test_capi/test_list.py

@@ -1,6 +1,7 @@
 import threading
 import unittest
 from collections import UserList
+from test import support
 from test.support import import_helper
 from test.support import threading_helper
 
@@ -11,6 +12,7 @@
 NULL = None
 PY_SSIZE_T_MIN = _testcapi.PY_SSIZE_T_MIN
 PY_SSIZE_T_MAX = _testcapi.PY_SSIZE_T_MAX
+LIST_SMALL_ALLOCATED = _testcapi.LIST_SMALL_ALLOCATED
 
 class ListSubclass(list):
     pass
@@ -360,7 +362,7 @@ def test_uninitialized_list_repr(self):
 
 
 class ListWriterTests(unittest.TestCase):
-    def create_writer(self, size=0):
+    def create_writer(self, size):
         # Test PyListWriter_Create()
         return _testcapi.PyListWriter(size)
 
@@ -385,6 +387,16 @@ def check_append(self, method):
                     append(item)
                 self.list_equal(writer.finish(), items)
 
+        # test MemoryError
+        writer = self.create_writer(0)
+        append = getattr(writer, method)
+        with self.assertRaises(MemoryError):
+            with support.inject_memory_error_cm():
+                append("alloc failed")
+        for item in items:
+            append(item)
+        self.list_equal(writer.finish(), items)
+
     def test_append(self):
         # Test PyListWriter_Append()
         self.check_append('append')
@@ -402,6 +414,16 @@ def check_extend(self, method):
                 extend(items)
                 self.list_equal(writer.finish(), items)
 
+        # test MemoryError
+        writer = self.create_writer(0)
+        extend = getattr(writer, method)
+        items_error = ("alloc failed",)
+        with self.assertRaises(MemoryError):
+            with support.inject_memory_error_cm():
+                extend(items_error)
+        extend(items)
+        self.list_equal(writer.finish(), items)
+
     def test_extend(self):
         # Test PyListWriter_Extend()
         self.check_extend('extend')
@@ -414,27 +436,94 @@ def test_extend_array_steal(self):
         # Test PyListWriter_ExtendArraySteal()
         self.check_extend('extend_array_steal')
 
+    def test_empty(self):
+        list_get_allocated = _testcapi._list_get_allocated
+        empty_tuple = ()
+
+        for size in (0, LIST_SMALL_ALLOCATED, 1024):
+            with self.subTest(size=size):
+                writer = self.create_writer(size)
+                writer.extend(empty_tuple)
+                writer.extend_array(empty_tuple)
+                writer.extend_array_steal(empty_tuple)
+                result = writer.finish()
+                self.list_equal(result, [])
+
+                if size == 0 or size > LIST_SMALL_ALLOCATED:
+                    expected = 0
+                else:
+                    expected = size
+                self.assertEqual(list_get_allocated(result), expected)
+
+    def test_finish_shrink(self):
+        # PyListWriter_Finish() shrinks the list if its allocated size
+        # is way too big
+        list_get_allocated = _testcapi._list_get_allocated
+
+        def compute_allocated(size):
+            return (size + (size >> 3) + 6) & ~3
+
+        def check_allocated(items, alloc, expected):
+            writer = self.create_writer(alloc)
+            # Don't use PyListWriter_Extend() since it can shrink the list
+            for item in items:
+                writer.append(item)
+            result = writer.finish()
+            self.list_equal(result, list(items))
+            self.assertEqual(list_get_allocated(result), expected)
+
+        # small list
+        items = [1, 2, 3]
+        check_allocated(items,
+                        1024,
+                        compute_allocated(len(items)))
+
+        alloc = LIST_SMALL_ALLOCATED
+        check_allocated(items, alloc, alloc)
+
+        # larger list
+        items = tuple(range(100))
+        check_allocated(items,
+                        (len(items) + 1) * 2,
+                        compute_allocated(len(items)))
+
+        alloc = len(items) * 2
+        check_allocated(items, alloc, alloc)
+
+        # do not shrink on MemoryError
+        alloc = 1024
+        writer = self.create_writer(alloc)
+        writer.append(1)
+        with support.inject_memory_error_cm():
+            # list_resize() never fails when shrinking
+            result = writer.finish()
+        self.list_equal(result, [1])
+        self.assertEqual(list_get_allocated(result), alloc)
+
     def test_discard(self):
         # Test PyListWriter_Discard(): must not leak memory
-        writer = self.create_writer()
+        writer = self.create_writer(0)
         writer.extend([2, 5, 7])
         writer.discard()
 
     @threading_helper.requires_working_threading()
     def test_thread(self):
         # PyListWriter can be used by multiple threads: it's up to the caller
         # to implement a lock to prevent concurrent accesses.
-        writer = self.create_writer()
+        writer = self.create_writer(0)
 
         def thread_func(writer):
-            writer.append('Hello')
-            writer.extend(['World!'])
+            writer.append('a')
+            writer.append_steal('b')
+            writer.extend(('c',))
+            writer.extend_array(('d',))
+            writer.extend_array_steal(('e',))
 
         thread = threading.Thread(target=thread_func, args=(writer,))
         thread.start()
         threading_helper.join_thread(thread)
 
-        self.list_equal(writer.finish(), ['Hello', 'World!'])
+        self.list_equal(writer.finish(), list('abcde'))
 
 
 if __name__ == "__main__":

Modules/_testcapi/list.c

@@ -1,6 +1,11 @@
+// Use pycore_list.h
+#define PYTESTCAPI_NEED_INTERNAL_API
+
 #include "parts.h"
 #include "util.h"
 
+#include "pycore_list.h"          // _Py_LIST_SMALL_ALLOCATED
+
 
 static PyObject *
 list_get_size(PyObject *Py_UNUSED(module), PyObject *obj)
@@ -103,6 +108,15 @@ test_list_api(PyObject *self, PyObject *Py_UNUSED(ignored))
 }
 
 
+static PyObject *
+_list_get_allocated(PyObject *Py_UNUSED(module), PyObject *obj)
+{
+    NULLABLE(obj);
+    PyListObject *list = (PyListObject*)obj;
+    return PyLong_FromSsize_t(list->allocated);
+}
+
+
 typedef struct {
     PyObject_HEAD
     PyListWriter *writer;
@@ -338,6 +352,7 @@ static PyMethodDef test_methods[] = {
     {"list_clear", list_clear, METH_O},
     {"list_extend", list_extend, METH_VARARGS},
     {"test_list_api", test_list_api, METH_NOARGS},
+    {"_list_get_allocated", _list_get_allocated, METH_O},
     {NULL},
 };
 
@@ -358,5 +373,9 @@ _PyTestCapi_Init_List(PyObject *m)
     }
     Py_DECREF(writer_type);
 
+    if (PyModule_AddIntConstant(m, "LIST_SMALL_ALLOCATED", _Py_LIST_SMALL_ALLOCATED) < 0) {
+        return -1;
+    }
+
     return 0;
 }

Objects/listobject.c

@@ -91,7 +91,7 @@ ensure_shared_on_resize(PyListObject *self)
 #endif
 }
 
-#define LIST_SMALL_ALLOCATED 32
+#define LIST_SMALL_ALLOCATED _Py_LIST_SMALL_ALLOCATED
 
 static int py_list_resize(PyListObject *self, Py_ssize_t newsize);
 
@@ -4419,7 +4419,10 @@ static int
 listwriter_check_consistency(PyListWriter *writer)
 {
     assert(writer != NULL);
-    PyObject *list = _PyObject_CAST(writer);
+    assert(PyList_CheckExact(_PyObject_CAST(writer)));
+    PyListObject *list = (PyListObject*)writer;
+    assert(Py_SIZE(list) >= 0);
+    assert(Py_SIZE(list) <= list->allocated);
     // Don't test _PyObject_IsUniquelyReferenced() since a list writer can
     // be shared between multiple threads with external locking.
     assert(Py_REFCNT(list) == 1);
@@ -4479,6 +4482,22 @@ PyListWriter_Finish(PyListWriter *writer)
     assert(listwriter_check_consistency(writer));
     PyListObject *list = _PyList_CAST(writer);
 
+    // Shrink the list if the allocated size is greater than the double size.
+    // For example, if PyListWriter_Create() argument overestimated the final
+    // size.
+    if (Py_SIZE(list) < (list->allocated >> 1)
+        && list->allocated > LIST_SMALL_ALLOCATED)
+    {
+        if (list_resize(list, Py_SIZE(list)) < 0) {
+            // This case cannot happen: list_resize() never fails when
+            // shrinking
+            PyListWriter_Discard(writer);
+            return NULL;
+        }
+        // Note: cannot test that list->allocated changed in debug mode, since
+        // it's left unchanged on memory allocation failure.
+    }
+
 #ifdef Py_DEBUG
     // Make sure that all items have been initialized
     Py_ssize_t size = Py_SIZE(list);
@@ -4534,14 +4553,15 @@ PyListWriter_Append(PyListWriter *writer, PyObject *item)
 }
 
 
+// Similar to _list_extend() but don't lock the writer list.
 // If you update this function, update also _list_extend().
 int
 PyListWriter_Extend(PyListWriter *writer, PyObject *iterable)
 {
     assert(listwriter_check_consistency(writer));
-    assert(_PyObject_CAST(writer) != iterable);
-
     PyListObject *self = _PyList_CAST(writer);
+    assert(_PyObject_CAST(self) != iterable);
+
     int res;
     if (PyList_CheckExact(iterable)) {
         Py_BEGIN_CRITICAL_SECTION(iterable);