ENH: Add PyUnicode_FromUCS4 to ucsnarrow.

mwcraig · Apr 21, 2012 · 321a014 · 321a014
1 parent 8330e6f
commit 321a014
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 0 deletions.
diff --git a/numpy/core/src/multiarray/ucsnarrow.c b/numpy/core/src/multiarray/ucsnarrow.c
@@ -86,6 +86,90 @@ PyUCS2Buffer_AsUCS4(Py_UNICODE *ucs2, npy_ucs4 *ucs4, int ucs2len, int ucs4len)
     return numchars;
 }
 
+/*
+ * Returns a PyUnicodeObject initialized from a buffer containing
+ * UCS4 unicode.
+ *
+ * Parameters
+ * ----------
+ *  src: char *
+ *      Pointer to buffer containing UCS4 unicode.
+ *  size: Py_ssize_t
+ *      Size of buffer in bytes.
+ *  swap: int
+ *      If true, the data will be swapped.
+ *  align: int
+ *      If true, the data will be aligned.
+ *
+ * Returns
+ * -------
+ * new_reference: PyUnicodeObject
+ */ 
+NPY_NO_EXPORT PyUnicodeObject *
+PyUnicode_FromUCS4(char *src, Py_ssize_t size, int swap, int align)
+{
+    Py_ssize_t ucs4len = size / sizeof(npy_ucs4);
+    npy_ucs4 *buf = (npy_ucs4 *)src;
+    int alloc = 0;
+    PyUnicodeObject *ret;
+
+    /* swap and align if needed */
+    if (swap || align) {
+        buf = (npy_ucs4 *)malloc(size);
+        if (buf == NULL) {
+            PyErr_NoMemory();
+            goto fail;
+        }
+        alloc = 1;
+        memcpy(buf, src, size);
+        if (swap) {
+            byte_swap_vector(buf, ucs4len, sizeof(npy_ucs4));
+        }
+    }
+
+    /* trim trailing zeros */
+    while (ucs4len > 0 && buf[ucs4len - 1] == 0) {
+        ucs4len--;
+    }
+
+    /* produce PyUnicode object */
+#ifdef Py_UNICODE_WIDE
+    {
+        ret = (PyUnicodeObject *)PyUnicode_FromUnicode(buf, (Py_ssize_t) ucs4len);
+        if (ret == NULL) {
+            goto fail;
+        }
+    }
+#else
+    {
+        Py_ssize_t tmpsiz = 2 * sizeof(Py_UNICODE) * ucs4len;
+        Py_ssize_t ucs2len;
+        Py_UNICODE *tmp;
+
+        if ((tmp = (Py_UNICODE *)malloc(tmpsiz)) == NULL) {
+            PyErr_NoMemory();
+            goto fail;
+        }
+        ucs2len = PyUCS2Buffer_FromUCS4(tmp, buf, ucs4len);
+        ret = (PyUnicodeObject *)PyUnicode_FromUnicode(tmp, (Py_ssize_t) ucs2len);
+        free(tmp);
+        if (ret == NULL) {
+            goto fail;
+        }
+    }
+#endif
+
+    if (alloc) {
+        free(buf);
+    }
+    return ret;
+
+fail:
+    if (alloc) {
+        free(buf);
+    }
+    return NULL;   
+}
 
 NPY_NO_EXPORT PyObject *
 MyPyUnicode_New(int length)

diff --git a/numpy/core/src/multiarray/ucsnarrow.h b/numpy/core/src/multiarray/ucsnarrow.h
@@ -13,4 +13,7 @@ MyPyUnicode_New(int length);
 NPY_NO_EXPORT int
 MyPyUnicode_Resize(PyUnicodeObject *uni, int length);
 
+NPY_NO_EXPORT PyUnicodeObject *
+PyUnicode_FromUCS4(char *src, Py_ssize_t size, int swap, int align);
+
 #endif