gh-158585: Adjust _PyBytes_FormatEx() allocation

Count properly the number of preallocated bytes when growing the
buffer.

Author: vstinner

Lib/test/test_bytes.py

@@ -894,7 +894,7 @@ def __int__(self):
     def test_memory_leak_gh_140939(self):
         # gh-140939: MemoryError is raised without leaking
         _testcapi = import_helper.import_module('_testcapi')
-        with self.assertRaises(MemoryError):
+        with self.assertRaises((MemoryError, OverflowError)):
             b = self.type2test(b'%*b')
             b % (_testcapi.PY_SSIZE_T_MAX, b'abc')
 

Objects/bytesobject.c

@@ -747,6 +747,7 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
             char onechar; /* For byte_converter() */
             Py_ssize_t alloc;
 
+            const char *start = fmt;
             fmt++;
             if (*fmt == '%') {
                 *res++ = '%';
@@ -1108,9 +1109,10 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
             alloc = width;
             if (sign != 0 && len == width)
                 alloc++;
-            /* 2: size preallocated for %s */
-            if (alloc > 2) {
-                res = PyBytesWriter_GrowAndUpdatePointer(writer, alloc - 2, res);
+            /* size preallocated for the format */
+            Py_ssize_t prealloc = fmt - start;
+            if (alloc > prealloc) {
+                res = PyBytesWriter_GrowAndUpdatePointer(writer, alloc - prealloc, res);
                 if (res == NULL) {
                     Py_XDECREF(temp);
                     goto error;