Add PyBytesWriter_Format()

Author: vstinner

Include/cpython/bytesobject.h

@@ -70,6 +70,10 @@ PyAPI_FUNC(int) PyBytesWriter_WriteBytes(
     PyBytesWriter *writer,
     const void *bytes,
     Py_ssize_t size);
+PyAPI_FUNC(int) PyBytesWriter_Format(
+    PyBytesWriter *writer,
+    const char *format,
+    ...);
 
 PyAPI_FUNC(int) PyBytesWriter_Resize(
     PyBytesWriter *writer,

Lib/test/test_capi/test_bytes.py

@@ -356,11 +356,16 @@ def test_resize(self):
         writer.resize(len(b'number=123456'), b'456')
         self.assertEqual(writer.finish(), b'number=123456')
 
-    # def test_format(self):
-    #     # Test PyBytesWriter_Format()
-    #     writer = self.create_writer()
-    #     writer.format_i(123456)
-    #     self.assertEqual(writer.finish(), b'123456')
+    def test_format_i(self):
+        # Test PyBytesWriter_Format()
+        writer = self.create_writer()
+        writer.format_i(b'x=%i', 123456)
+        self.assertEqual(writer.finish(), b'x=123456')
+
+        writer = self.create_writer()
+        writer.format_i(b'x=%i, ', 123)
+        writer.format_i(b'y=%i', 456)
+        self.assertEqual(writer.finish(), b'x=123, y=456')
 
 
 if __name__ == "__main__":

Modules/_testcapi/bytes.c

@@ -151,6 +151,27 @@ writer_write_bytes(PyObject *self_raw, PyObject *args)
 }
 
 
+static PyObject*
+writer_format_i(PyObject *self_raw, PyObject *args)
+{
+    WriterObject *self = (WriterObject *)self_raw;
+    if (writer_check(self) < 0) {
+        return NULL;
+    }
+
+    char *format;
+    int value;
+    if (!PyArg_ParseTuple(args, "yi", &format, &value)) {
+        return NULL;
+    }
+
+    if (PyBytesWriter_Format(self->writer, format, value) < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyObject*
 writer_resize(PyObject *self_raw, PyObject *args)
 {
@@ -242,6 +263,7 @@ writer_finish_with_size(PyObject *self_raw, PyObject *args)
 
 static PyMethodDef writer_methods[] = {
     {"write_bytes", _PyCFunction_CAST(writer_write_bytes), METH_VARARGS},
+    {"format_i", _PyCFunction_CAST(writer_format_i), METH_VARARGS},
     {"resize", _PyCFunction_CAST(writer_resize), METH_VARARGS},
     {"get_size", _PyCFunction_CAST(writer_get_size), METH_NOARGS},
     {"get_allocated", _PyCFunction_CAST(writer_get_allocated), METH_NOARGS},

Objects/bytesobject.c

@@ -196,8 +196,10 @@ PyBytes_FromString(const char *str)
     return (PyObject *) op;
 }
 
-PyObject *
-PyBytes_FromFormatV(const char *format, va_list vargs)
+
+static char*
+bytes_fromformat(PyBytesWriter *writer, Py_ssize_t writer_pos,
+                 const char *format, va_list vargs)
 {
     const char *f;
     const char *p;
@@ -213,12 +215,8 @@ PyBytes_FromFormatV(const char *format, va_list vargs)
        "0xffffffffffffffff\0" (19 bytes). */
     char buffer[21];
 
-    Py_ssize_t alloc = strlen(format);
-    PyBytesWriter *writer = PyBytesWriter_Create(alloc);
-    if (writer == NULL) {
-        return NULL;
-    }
-    char *s = PyBytesWriter_GetData(writer);
+    char *s = PyBytesWriter_GetData(writer) + writer_pos;
+    Py_ssize_t alloc = PyBytesWriter_GetSize(writer);
 
 #define WRITE_BYTES_LEN(str, len_expr) \
     do { \
@@ -366,20 +364,39 @@ PyBytes_FromFormatV(const char *format, va_list vargs)
         default:
             /* invalid format string: copy unformatted string and exit */
             WRITE_BYTES(p);
-            return PyBytesWriter_FinishWithEndPointer(writer, s);
+            return s;
         }
     }
 
 #undef WRITE_BYTES
 #undef WRITE_BYTES_LEN
 
-    return PyBytesWriter_FinishWithEndPointer(writer, s);
+    return s;
 
  error:
-    PyBytesWriter_Discard(writer);
     return NULL;
 }
 
+
+PyObject *
+PyBytes_FromFormatV(const char *format, va_list vargs)
+{
+    Py_ssize_t alloc = strlen(format);
+    PyBytesWriter *writer = PyBytesWriter_Create(alloc);
+    if (writer == NULL) {
+        return NULL;
+    }
+
+    char *s = bytes_fromformat(writer, 0, format, vargs);
+    if (s == NULL) {
+        PyBytesWriter_Discard(writer);
+        return NULL;
+    }
+
+    return PyBytesWriter_FinishWithEndPointer(writer, s);
+}
+
+
 PyObject *
 PyBytes_FromFormat(const char *format, ...)
 {
@@ -392,6 +409,7 @@ PyBytes_FromFormat(const char *format, ...)
     return ret;
 }
 
+
 /* Helpers for formatstring */
 
 Py_LOCAL_INLINE(PyObject *)
@@ -3949,3 +3967,28 @@ PyBytesWriter_WriteBytes(PyBytesWriter *writer,
     memcpy(buf + pos, bytes, size);
     return 0;
 }
+
+
+int
+PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
+{
+    Py_ssize_t pos = writer->size;
+    Py_ssize_t format_len = strlen(format);
+    if (format_len > PY_SSIZE_T_MAX - pos) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    Py_ssize_t alloc = pos + format_len;
+
+    if (PyBytesWriter_Resize(writer, alloc) < 0) {
+        return -1;
+    }
+
+    va_list vargs;
+    va_start(vargs, format);
+    char *buf = bytes_fromformat(writer, pos, format, vargs);
+    va_end(vargs);
+
+    Py_ssize_t size = buf - byteswriter_data(writer);
+    return PyBytesWriter_Resize(writer, size);
+}