packer: Use PyUnicode_AsUTF8AndSize() for utf-8 (#272)

1 files changed, 47 insertions, 0 deletions

diff --git a/msgpack/pack.h b/msgpack/pack.h
index 3bc21ea..4f3ce1d 100644
--- a/msgpack/pack.h
+++ b/msgpack/pack.h
@@ -67,6 +67,53 @@ static inline int msgpack_pack_write(msgpack_packer* pk, const char *data, size_
 
 #include "pack_template.h"
 
+// return -2 when o is too long
+static inline int
+msgpack_pack_unicode(msgpack_packer *pk, PyObject *o, long long limit)
+{
+#if PY_MAJOR_VERSION >= 3
+    assert(PyUnicode_Check(o));
+
+    Py_ssize_t len;
+    const char* buf = PyUnicode_AsUTF8AndSize(o, &len);
+    if (buf == NULL)
+        return -1;
+
+    if (len > limit) {
+        return -2;
+    }
+
+    int ret = msgpack_pack_raw(pk, len);
+    if (ret) return ret;
+
+    return msgpack_pack_raw_body(pk, buf, len);
+#else
+    PyObject *bytes;
+    Py_ssize_t len;
+    int ret;
+
+    // py2
+    bytes = PyUnicode_AsUTF8String(o);
+    if (bytes == NULL)
+        return -1;
+
+    len = PyString_GET_SIZE(bytes);
+    if (len > limit) {
+        Py_DECREF(bytes);
+        return -2;
+    }
+
+    ret = msgpack_pack_raw(pk, len);
+    if (ret) {
+        Py_DECREF(bytes);
+        return -1;
+    }
+    ret = msgpack_pack_raw_body(pk, PyString_AS_STRING(bytes), len);
+    Py_DECREF(bytes);
+    return ret;
+#endif
+}
+
 #ifdef __cplusplus
 }
 #endif