Optimize bytes_fromformat()

Author: vstinner

Lib/test/test_capi/test_bytes.py

@@ -291,15 +291,25 @@ def test_join(self):
 
 
 class PyBytesWriterTest(unittest.TestCase):
-    def create_writer(self, alloc):
-        return _testcapi.PyBytesWriter(alloc)
+    def create_writer(self, alloc, string=b''):
+        return _testcapi.PyBytesWriter(alloc, string)
 
     def test_empty(self):
         # Test PyBytesWriter_Create()
         writer = self.create_writer(0)
         self.assertEqual(writer.get_remaining(), 0)
         self.assertEqual(writer.finish(), b'')
 
+    def test_abc(self):
+        # Test PyBytesWriter_Create()
+        writer = self.create_writer(3, b'abc')
+        self.assertEqual(writer.get_remaining(), 0)
+        self.assertEqual(writer.finish(), b'abc')
+
+        writer = self.create_writer(10, b'abc')
+        self.assertEqual(writer.get_remaining(), 7)
+        self.assertEqual(writer.finish(), b'abc')
+
     def test_write_bytes(self):
         # Test PyBytesWriter_WriteBytes()
 
@@ -352,6 +362,7 @@ def test_format(self):
         # Test PyBytesWriter_Format()
         writer = self.create_writer(0)
         writer.format_i(123456)
+        self.assertEqual(writer.get_remaining(), 0)
         self.assertEqual(writer.finish(), b'123456')
 
 

Modules/_testcapi/bytes.c

@@ -88,14 +88,20 @@ writer_init(PyObject *self_raw, PyObject *args, PyObject *kwargs)
     }
 
     Py_ssize_t alloc;
-    if (!PyArg_ParseTuple(args, "n", &alloc)) {
+    char *str;
+    Py_ssize_t str_size;
+    if (!PyArg_ParseTuple(args, "ny#", &alloc, &str, &str_size)) {
         return -1;
     }
 
     self->buf = PyBytesWriter_Create(&self->writer, alloc);
     if (self->buf == NULL) {
         return -1;
     }
+
+    memcpy(self->buf, str, str_size);
+    self->buf += str_size;
+
     return 0;
 }
 

Objects/bytesobject.c

@@ -215,10 +215,15 @@ bytes_fromformat(PyBytesWriter *writer, char *s,
 
 #define WRITE_BYTES(str) \
     do { \
-        s = PyBytesWriter_WriteBytes(writer, s, (str), strlen(str)); \
-        if (s == NULL) { \
-            goto error; \
+        Py_ssize_t len = strlen(str); \
+        if (len > 2) { \
+            s = PyBytesWriter_Extend(writer, s, len - 2); \
+            if (s == NULL) { \
+                goto error; \
+            } \
         } \
+        memcpy(s, (str), len); \
+        s += len; \
     } while (0)
 
     for (f = format; *f; f++) {
@@ -355,7 +360,10 @@ bytes_fromformat(PyBytesWriter *writer, char *s,
 
         default:
             /* invalid format string: copy unformatted string and exit */
-            WRITE_BYTES(p);
+            s = PyBytesWriter_WriteBytes(writer, s, p, strlen(p));
+            if (s == NULL) {
+                goto error;
+            }
             goto done;
         }
     }