Issue #8870: PyUnicode_AsWideCharString() doesn't count the trailing nul character

And write unit tests for PyUnicode_AsWideChar() and PyUnicode_AsWideCharString().

4 files changed, 105 insertions, 11 deletions

diff --git a/Doc/c-api/unicode.rst b/Doc/c-api/unicode.rst
index 50a1254..45c46aa 100644
--- a/Doc/c-api/unicode.rst
+++ b/Doc/c-api/unicode.rst
@@ -466,7 +466,8 @@ wchar_t support for platforms which support it:
 
    Convert the Unicode object to a wide character string. The output string
    always ends with a nul character. If *size* is not *NULL*, write the number
-   of wide characters (including the nul character) into *\*size*.
+   of wide characters (excluding the trailing 0-termination character) into
+   *\*size*.
 
    Returns a buffer allocated by :cfunc:`PyMem_Alloc` (use :cfunc:`PyMem_Free`
    to free it) on success. On error, returns *NULL*, *\*size* is undefined and
diff --git a/Lib/test/test_unicode.py b/Lib/test/test_unicode.py
index f6c38dd..2ac79fb 100644
--- a/Lib/test/test_unicode.py
+++ b/Lib/test/test_unicode.py
@@ -1394,6 +1394,45 @@ class UnicodeTest(string_tests.CommonTest,
             'string, got a non-ASCII byte: 0xe9$',
             format_unicode, b'unicode\xe9=%s', 'ascii')
 
+    # Test PyUnicode_AsWideChar()
+    def test_aswidechar(self):
+        from _testcapi import test_aswidechar
+        from ctypes import c_wchar, sizeof
+
+        wchar, size = test_aswidechar('abcdef', 2)
+        self.assertEquals(size, 2)
+        self.assertEquals(wchar, 'ab')
+
+        wchar, size = test_aswidechar('abc', 3)
+        self.assertEquals(size, 3)
+        self.assertEquals(wchar, 'abc')
+
+        wchar, size = test_aswidechar('abc', 4)
+        self.assertEquals(size, 3)
+        self.assertEquals(wchar, 'abc\0')
+
+        wchar, size = test_aswidechar('abc', 10)
+        self.assertEquals(size, 3)
+        self.assertEquals(wchar, 'abc\0')
+
+        wchar, size = test_aswidechar('abc\0def', 20)
+        self.assertEquals(size, 7)
+        self.assertEquals(wchar, 'abc\0def\0')
+
+    # Test PyUnicode_AsWideCharString()
+    def test_aswidecharstring(self):
+        from _testcapi import test_aswidecharstring
+        from ctypes import c_wchar, sizeof
+
+        wchar, size = test_aswidecharstring('abc')
+        self.assertEquals(size, 3)
+        self.assertEquals(wchar, 'abc\0')
+
+        wchar, size = test_aswidecharstring('abc\0def')
+        self.assertEquals(size, 7)
+        self.assertEquals(wchar, 'abc\0def\0')
+
+
 def test_main():
     support.run_unittest(__name__)
 
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 20887b1..d2d3b19 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -1386,6 +1386,58 @@ test_widechar(PyObject *self)
 }
 
 static PyObject *
+test_aswidechar(PyObject *self, PyObject *args)
+{
+    PyObject *unicode, *result;
+    Py_ssize_t buflen, size;
+    wchar_t *buffer;
+
+    if (!PyArg_ParseTuple(args, "Un", &unicode, &buflen))
+        return NULL;
+    buffer = PyMem_Malloc(buflen * sizeof(wchar_t));
+    if (buffer == NULL)
+        return PyErr_NoMemory();
+
+    size = PyUnicode_AsWideChar((PyUnicodeObject*)unicode, buffer, buflen);
+    if (size == -1) {
+        PyMem_Free(buffer);
+        return NULL;
+    }
+
+    if (size < buflen)
+        buflen = size + 1;
+    else
+        buflen = size;
+    result = PyUnicode_FromWideChar(buffer, buflen);
+    PyMem_Free(buffer);
+    if (result == NULL)
+        return NULL;
+
+    return Py_BuildValue("(Nn)", result, size);
+}
+
+static PyObject *
+test_aswidecharstring(PyObject *self, PyObject *args)
+{
+    PyObject *unicode, *result;
+    Py_ssize_t size;
+    wchar_t *buffer;
+
+    if (!PyArg_ParseTuple(args, "U", &unicode))
+        return NULL;
+
+    buffer = PyUnicode_AsWideCharString((PyUnicodeObject*)unicode, &size);
+    if (buffer == NULL)
+        return NULL;
+
+    result = PyUnicode_FromWideChar(buffer, size + 1);
+    PyMem_Free(buffer);
+    if (result == NULL)
+        return NULL;
+    return Py_BuildValue("(Nn)", result, size);
+}
+
+static PyObject *
 getargs_w_star(PyObject *self, PyObject *args)
 {
     Py_buffer buffer;
@@ -2262,28 +2314,30 @@ static PyMethodDef TestMethods[] = {
     {"getargs_Z_hash",          getargs_Z_hash,                  METH_VARARGS},
     {"getargs_w_star",          getargs_w_star,                  METH_VARARGS},
     {"codec_incrementalencoder",
-     (PyCFunction)codec_incrementalencoder,      METH_VARARGS},
+     (PyCFunction)codec_incrementalencoder,                      METH_VARARGS},
     {"codec_incrementaldecoder",
-     (PyCFunction)codec_incrementaldecoder,      METH_VARARGS},
+     (PyCFunction)codec_incrementaldecoder,                      METH_VARARGS},
     {"test_s_code",             (PyCFunction)test_s_code,        METH_NOARGS},
     {"test_u_code",             (PyCFunction)test_u_code,        METH_NOARGS},
     {"test_Z_code",             (PyCFunction)test_Z_code,        METH_NOARGS},
     {"test_widechar",           (PyCFunction)test_widechar,      METH_NOARGS},
+    {"test_aswidechar",         test_aswidechar,                 METH_VARARGS},
+    {"test_aswidecharstring",   test_aswidecharstring,           METH_VARARGS},
 #ifdef WITH_THREAD
-    {"_test_thread_state",  test_thread_state,                   METH_VARARGS},
+    {"_test_thread_state",      test_thread_state,               METH_VARARGS},
     {"_pending_threadfunc",     pending_threadfunc,              METH_VARARGS},
 #endif
 #ifdef HAVE_GETTIMEOFDAY
-    {"profile_int",             profile_int,                    METH_NOARGS},
+    {"profile_int",             profile_int,                     METH_NOARGS},
 #endif
-    {"traceback_print", traceback_print,                 METH_VARARGS},
-    {"exception_print", exception_print,                 METH_VARARGS},
-    {"argparsing",     argparsing, METH_VARARGS},
-    {"code_newempty", code_newempty,                     METH_VARARGS},
+    {"traceback_print",         traceback_print,                 METH_VARARGS},
+    {"exception_print",         exception_print,                 METH_VARARGS},
+    {"argparsing",              argparsing,                      METH_VARARGS},
+    {"code_newempty",           code_newempty,                   METH_VARARGS},
     {"make_exception_with_doc", (PyCFunction)make_exception_with_doc,
      METH_VARARGS | METH_KEYWORDS},
     {"crash_no_current_thread", (PyCFunction)crash_no_current_thread, METH_NOARGS},
-    {"format_unicode", format_unicode, METH_VARARGS},
+    {"format_unicode",          format_unicode,                 METH_VARARGS},
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c
index 29404c3..1c083b2 100644
--- a/Objects/unicodeobject.c
+++ b/Objects/unicodeobject.c
@@ -1216,7 +1216,7 @@ PyUnicode_AsWideCharString(PyUnicodeObject *unicode,
     }
     unicode_aswidechar(unicode, buffer, buflen);
     if (size)
-        *size = buflen;
+        *size = buflen - 1;
     return buffer;
 }