Add _PyBytesWriter_WriteBytes() to factorize the code

4 files changed, 35 insertions, 16 deletions

diff --git a/Include/bytesobject.h b/Include/bytesobject.h
index 6cd5a34..2c4c4c4 100644
--- a/Include/bytesobject.h
+++ b/Include/bytesobject.h
@@ -174,6 +174,13 @@ PyAPI_FUNC(char*) _PyBytesWriter_Alloc(_PyBytesWriter *writer,
 PyAPI_FUNC(char*) _PyBytesWriter_Prepare(_PyBytesWriter *writer,
     char *str,
     Py_ssize_t size);
+
+/* Write bytes.
+   Raise an exception and return NULL on error. */
+PyAPI_FUNC(char*) _PyBytesWriter_WriteBytes(_PyBytesWriter *writer,
+    char *str,
+    char *bytes,
+    Py_ssize_t size);
 #endif   /* Py_LIMITED_API */
 
 #ifdef __cplusplus
diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c
index 075edf8..3aa905c 100644
--- a/Objects/bytesobject.c
+++ b/Objects/bytesobject.c
@@ -3995,3 +3995,17 @@ _PyBytesWriter_Finish(_PyBytesWriter *writer, char *str)
 
     return result;
 }
+
+char*
+_PyBytesWriter_WriteBytes(_PyBytesWriter *writer, char *str,
+                          char *bytes, Py_ssize_t size)
+{
+    str = _PyBytesWriter_Prepare(writer, str, size);
+    if (str == NULL)
+        return NULL;
+
+    Py_MEMCPY(str, bytes, size);
+    str += size;
+
+    return str;
+}
diff --git a/Objects/stringlib/codecs.h b/Objects/stringlib/codecs.h
index 6842f67..7e8d928 100644
--- a/Objects/stringlib/codecs.h
+++ b/Objects/stringlib/codecs.h
@@ -388,24 +388,24 @@ STRINGLIB(utf8_encoder)(PyObject *unicode,
                 /* substract preallocated bytes */
                 writer.min_size -= max_char_size;
 
-                if (PyBytes_Check(rep))
-                    repsize = PyBytes_GET_SIZE(rep);
-                else
-                    repsize = PyUnicode_GET_LENGTH(rep);
-
-                p = _PyBytesWriter_Prepare(&writer, p, repsize);
-                if (p == NULL)
-                    goto error;
-
                 if (PyBytes_Check(rep)) {
-                    memcpy(p, PyBytes_AS_STRING(rep), repsize);
-                    p += repsize;
+                    p = _PyBytesWriter_WriteBytes(&writer, p,
+                                                  PyBytes_AS_STRING(rep),
+                                                  PyBytes_GET_SIZE(rep));
+                    if (p == NULL)
+                        goto error;
                 }
                 else {
                     /* rep is unicode */
                     if (PyUnicode_READY(rep) < 0)
                         goto error;
 
+                    repsize = PyUnicode_GET_LENGTH(rep);
+
+                    p = _PyBytesWriter_Prepare(&writer, p, repsize);
+                    if (p == NULL)
+                        goto error;
+
                     if (!PyUnicode_IS_ASCII(rep)) {
                         raise_encode_exception(&exc, "utf-8",
                                                unicode,
diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c
index 0bcacd8..23b8cc7 100644
--- a/Objects/unicodeobject.c
+++ b/Objects/unicodeobject.c
@@ -6706,14 +6706,12 @@ unicode_encode_ucs1(PyObject *unicode,
 
                 if (PyBytes_Check(repunicode)) {
                     /* Directly copy bytes result to output. */
-                    repsize = PyBytes_Size(repunicode);
-
-                    str = _PyBytesWriter_Prepare(&writer, str, repsize);
+                    str = _PyBytesWriter_WriteBytes(&writer, str,
+                                                    PyBytes_AS_STRING(repunicode),
+                                                    PyBytes_GET_SIZE(repunicode));
                     if (str == NULL)
                         goto onError;
 
-                    memcpy(str, PyBytes_AsString(repunicode), repsize);
-                    str += repsize;
                     pos = newpos;
                     Py_DECREF(repunicode);
                     break;