in-progress

1 files changed, 43 insertions, 4 deletions

diff --git a/c/wchar_helper.h b/c/wchar_helper.h
index 9e88e45..f17bef3 100644
--- a/c/wchar_helper.h
+++ b/c/wchar_helper.h
@@ -63,6 +63,11 @@ _my_PyUnicode_FromWideChar(register const wchar_t *w,
 #endif
 
 
+#define IS_SURROGATE(u)   (0xD800 <= (u)[0] && (u)[0] <= 0xDBFF &&   \
+                           0xDC00 <= (u)[1] && (u)[1] <= 0xDFFF)
+#define AS_SURROGATE(u)   (0x10000 + (((u)[0] - 0xD800) << 10) +     \
+                                     ((u)[1] - 0xDC00))
+
 static int _my_PyUnicode_AsSingleWideChar(PyObject *unicode, wchar_t *result)
 {
     Py_UNICODE *u = PyUnicode_AS_UNICODE(unicode);
@@ -71,12 +76,46 @@ static int _my_PyUnicode_AsSingleWideChar(PyObject *unicode, wchar_t *result)
         return 0;
     }
 #ifdef CONVERT_WCHAR_TO_SURROGATES
-    if (PyUnicode_GET_SIZE(unicode) == 2 &&
-            0xD800 <= u[0] && u[0] <= 0xDBFF &&
-            0xDC00 <= u[1] && u[1] <= 0xDFFF) {
-        *result = 0x10000 + ((u[0] - 0xD800) << 10) + (u[1] - 0xDC00);
+    if (PyUnicode_GET_SIZE(unicode) == 2 && IS_SURROGATE(u)) {
+        *result = AS_SURROGATE(u);
         return 0;
     }
 #endif
     return -1;
 }
+
+static Py_ssize_t _my_PyUnicode_SizeAsWideChar(PyObject *unicode)
+{
+    Py_ssize_t length = PyUnicode_GET_SIZE(unicode);
+    Py_ssize_t result = length;
+
+#ifdef CONVERT_WCHAR_TO_SURROGATES
+    Py_UNICODE *u = PyUnicode_AS_UNICODE(unicode);
+    Py_ssize_t i;
+
+    for (i=0; i<length-1; i++) {
+        if (IS_SURROGATE(u+i))
+            result--;
+    }
+#endif
+    return result;
+}
+
+static void _my_PyUnicode_AsWideChar(PyObject *unicode,
+                                     wchar_t *result,
+                                     Py_ssize_t resultlen)
+{
+    Py_UNICODE *u = PyUnicode_AS_UNICODE(unicode);
+    Py_ssize_t i;
+    for (i=0; i<resultlen; i++) {
+        wchar_t ordinal = *u;
+#ifdef CONVERT_WCHAR_TO_SURROGATES
+        if (IS_SURROGATE(u)) {
+            ordinal = AS_SURROGATE(u);
+            u++;
+        }
+#endif
+        result[i] = ordinal;
+        u++;
+    }
+}