Add output_length to _Py_EncodeLocaleEx()

Add output_length to _Py_EncodeLocaleEx(), _Py_EncodeUTF8Ex(),
encode_current_locale() and encode_ascii(). So
unicode_encode_locale() and wstr_to_utf8() can use the output_length,
instead of having to compute strlen().

* Add encode_current_locale_impl() to simplify encode_current_locale().
* _Py_EncodeLocaleEx() now sets error_pos and reason if it fails with
  -1 or -3.
* Add tests on Py_EncodeLocale() and Py_DecodeLocale() functions in
  test_codecs.
* Remove reason parameter of _Py_EncodeUTF8Ex(),
  encode_current_locale() and encode_ascii().

Author: vstinner

Include/internal/pycore_fileutils.h

@@ -36,6 +36,7 @@ PyAPI_FUNC(int) _Py_DecodeLocaleEx(
 PyAPI_FUNC(int) _Py_EncodeLocaleEx(
     const wchar_t *text,
     char **str,
+    size_t *output_length,
     size_t *error_pos,
     const char **reason,
     int current_locale,
@@ -201,8 +202,8 @@ extern int _Py_DecodeUTF8Ex(
 extern int _Py_EncodeUTF8Ex(
     const wchar_t *text,
     char **str,
+    size_t *output_length,
     size_t *error_pos,
-    const char **reason,
     int raw_malloc,
     _Py_error_handler errors);
 

Lib/test/test_codecs.py

@@ -4057,6 +4057,7 @@ def test_pickle(self):
                     pickle.dumps(sr, proto)
 
 
+@unittest.skipIf(_testlimitedcapi is None, 'need _testlimitedcapi module')
 @unittest.skipIf(_testinternalcapi is None, 'need _testinternalcapi module')
 class LocaleCodecTest(unittest.TestCase):
     """
@@ -4070,7 +4071,12 @@ class LocaleCodecTest(unittest.TestCase):
     BYTES_STRINGS = (b"blatin1:\xa7\xe9", b"b255:\xff")
     SURROGATES = "\uDC80\uDCFF"
 
-    def encode(self, text, errors="strict"):
+    def encode_locale(self, text):
+        # Test Py_EncodeLocale(): use the "surrogateescape" error handler
+        return _testlimitedcapi.encode_locale(text)
+
+    def encode_locale_ex(self, text, errors="strict"):
+        # Test _Py_EncodeLocaleEx()
         return _testinternalcapi.EncodeLocaleEx(text, 0, errors)
 
     def check_encode_strings(self, errors):
@@ -4079,12 +4085,30 @@ def check_encode_strings(self, errors):
                 try:
                     expected = text.encode(self.ENCODING, errors)
                 except UnicodeEncodeError:
+                    for error_pos in range(len(text)):
+                        try:
+                            text[error_pos].encode(self.ENCODING, errors)
+                        except UnicodeEncodeError:
+                            break
+                    else:
+                        self.fail("failed to compute error_pos")
+
+                    if errors == "surrogateescape":
+                        with self.assertRaises(ValueError) as cm:
+                            self.encode_locale(text)
+                        errmsg = str(cm.exception)
+                        self.assertRegex(errmsg, f"Py_EncodeLocale failed: error_pos={error_pos}")
+
                     with self.assertRaises(RuntimeError) as cm:
-                        self.encode(text, errors)
+                        self.encode_locale_ex(text, errors)
                     errmsg = str(cm.exception)
-                    self.assertRegex(errmsg, r"encode error: pos=[0-9]+, reason=")
+                    self.assertRegex(errmsg, f"encode error: pos={error_pos}, reason=encoding error")
                 else:
-                    encoded = self.encode(text, errors)
+                    if errors in ("strict", "surrogateescape"):
+                        encoded = self.encode_locale(text)
+                        self.assertEqual(encoded, expected)
+
+                    encoded = self.encode_locale_ex(text, errors)
                     self.assertEqual(encoded, expected)
 
     def test_encode_strict(self):
@@ -4095,7 +4119,7 @@ def test_encode_surrogateescape(self):
 
     def test_encode_surrogatepass(self):
         try:
-            self.encode('', 'surrogatepass')
+            self.encode_locale_ex('', 'surrogatepass')
         except ValueError as exc:
             if str(exc) == 'unsupported error handler':
                 self.skipTest(f"{self.ENCODING!r} encoder doesn't support "
@@ -4107,12 +4131,17 @@ def test_encode_surrogatepass(self):
 
     def test_encode_unsupported_error_handler(self):
         with self.assertRaises(ValueError) as cm:
-            self.encode('', 'backslashreplace')
+            self.encode_locale_ex('', 'backslashreplace')
         self.assertEqual(str(cm.exception), 'unsupported error handler')
 
-    def decode(self, encoded, errors="strict"):
+    def decode_locale_ex(self, encoded, errors="strict"):
+        # Test _Py_DecodeLocaleEx()
         return _testinternalcapi.DecodeLocaleEx(encoded, 0, errors)
 
+    def decode_locale(self, encoded):
+        # Test DecodeLocale(): use the "surrogateescape" error handler
+        return _testlimitedcapi.decode_locale(encoded)
+
     def check_decode_strings(self, errors):
         is_utf8 = (self.ENCODING == "utf-8")
         if is_utf8:
@@ -4139,12 +4168,20 @@ def check_decode_strings(self, errors):
                 try:
                     expected = encoded.decode(self.ENCODING, errors)
                 except UnicodeDecodeError:
+                    if errors == "surrogateescape":
+                        with self.assertRaises(ValueError):
+                            self.decode_locale(encoded)
+
                     with self.assertRaises(RuntimeError) as cm:
-                        self.decode(encoded, errors)
+                        self.decode_locale_ex(encoded, errors)
                     errmsg = str(cm.exception)
                     self.assertStartsWith(errmsg, "decode error: ")
                 else:
-                    decoded = self.decode(encoded, errors)
+                    if errors == ("strict", "surrogateescape"):
+                        decoded = self.decode_locale(encoded)
+                        self.assertEqual(decoded, expected)
+
+                    decoded = self.decode_locale_ex(encoded, errors)
                     self.assertEqual(decoded, expected)
 
     def test_decode_strict(self):
@@ -4155,7 +4192,7 @@ def test_decode_surrogateescape(self):
 
     def test_decode_surrogatepass(self):
         try:
-            self.decode(b'', 'surrogatepass')
+            self.decode_locale_ex(b'', 'surrogatepass')
         except ValueError as exc:
             if str(exc) == 'unsupported error handler':
                 self.skipTest(f"{self.ENCODING!r} decoder doesn't support "
@@ -4167,7 +4204,7 @@ def test_decode_surrogatepass(self):
 
     def test_decode_unsupported_error_handler(self):
         with self.assertRaises(ValueError) as cm:
-            self.decode(b'', 'backslashreplace')
+            self.decode_locale_ex(b'', 'backslashreplace')
         self.assertEqual(str(cm.exception), 'unsupported error handler')
 
 

Modules/_testinternalcapi.c

@@ -1046,55 +1046,72 @@ get_getpath_codeobject(PyObject *self, PyObject *Py_UNUSED(args)) {
 }
 
 
+// Test _Py_EncodeLocaleEx()
 static PyObject *
 encode_locale_ex(PyObject *self, PyObject *args)
 {
     PyObject *unicode;
     int current_locale = 0;
-    wchar_t *wstr;
     PyObject *res = NULL;
     const char *errors = NULL;
 
     if (!PyArg_ParseTuple(args, "U|is", &unicode, &current_locale, &errors)) {
         return NULL;
     }
-    wstr = PyUnicode_AsWideCharString(unicode, NULL);
+
+    wchar_t *wstr = PyUnicode_AsWideCharString(unicode, NULL);
     if (wstr == NULL) {
         return NULL;
     }
     _Py_error_handler error_handler = _Py_GetErrorHandler(errors);
 
     char *str = NULL;
-    size_t error_pos;
-    const char *reason = NULL;
+    size_t error_pos_canary = (size_t)-123;
+    size_t error_pos = error_pos_canary;
+    size_t output_length = (size_t)-123;
+    const char *reason_canary = "canary";
+    const char *reason = reason_canary;
     int ret = _Py_EncodeLocaleEx(wstr,
-                                 &str, &error_pos, &reason,
+                                 &str, &output_length, &error_pos, &reason,
                                  current_locale, error_handler);
     PyMem_Free(wstr);
 
     switch(ret) {
     case 0:
-        res = PyBytes_FromString(str);
+        res = PyBytes_FromStringAndSize(str, output_length);
         PyMem_RawFree(str);
         break;
     case -1:
+        assert(output_length == 0);
+        assert(error_pos == 0);
+        assert(reason == NULL);
         PyErr_NoMemory();
         break;
     case -2:
+        assert(output_length == 0);
+        assert(error_pos != error_pos_canary);
+        assert(reason != reason_canary);
         PyErr_Format(PyExc_RuntimeError, "encode error: pos=%zu, reason=%s",
                      error_pos, reason);
         break;
     case -3:
+        assert(output_length == 0);
+        assert(error_pos == 0);
+        assert(reason == NULL);
         PyErr_SetString(PyExc_ValueError, "unsupported error handler");
         break;
     default:
+        assert(output_length == 0);
+        assert(error_pos == 0);
+        assert(reason == NULL);
         PyErr_SetString(PyExc_ValueError, "unknown error code");
         break;
     }
     return res;
 }
 
 
+// Test _Py_DecodeLocaleEx()
 static PyObject *
 decode_locale_ex(PyObject *self, PyObject *args)
 {

Modules/_testlimitedcapi/codec.c

@@ -2,8 +2,8 @@
 #ifdef Py_GIL_DISABLED
 #  define Py_TARGET_ABI3T 0x030f0000
 #else
-   // Need limited C API version 3.5 for PyCodec_NameReplaceErrors()
-#  define Py_LIMITED_API 0x03050000
+   // Need limited C API version 3.13 for PyMem_RawFree()
+#  define Py_LIMITED_API 0x030d0000
 #endif
 
 #include "parts.h"
@@ -15,16 +15,80 @@ codec_namereplace_errors(PyObject *Py_UNUSED(module), PyObject *exc)
     return PyCodec_NameReplaceErrors(exc);
 }
 
+
+// Test Py_DecodeLocale()
+static PyObject *
+decode_locale(PyObject *Py_UNUSED(module), PyObject *arg)
+{
+    const char *str;
+    if (PyArg_Parse(arg, "y", &str) < 0) {
+        return NULL;
+    }
+
+    size_t wstr_len = (size_t)-123;
+    wchar_t *wstr = Py_DecodeLocale(str, &wstr_len);
+
+    if (str == NULL) {
+        if (wstr_len == (size_t)-1) {
+            PyErr_NoMemory();
+        }
+        else if (wstr_len == (size_t)-2) {
+            PyErr_SetString(PyExc_ValueError, "decode error");
+        }
+        else {
+            PyErr_Format(PyExc_SystemError,
+                         "unknown Py_DecodeLocale() return value: %zd",
+                         (Py_ssize_t)wstr_len);
+        }
+        return NULL;
+    }
+
+    PyObject *result = PyUnicode_FromWideChar(wstr, wstr_len);
+    PyMem_RawFree(wstr);
+    return result;
+}
+
+
+// Test Py_EncodeLocale()
+static PyObject *
+encode_locale(PyObject *Py_UNUSED(module), PyObject *arg)
+{
+    PyObject *unicode;
+    if (PyArg_Parse(arg, "U", &unicode) < 0) {
+        return NULL;
+    }
+
+    wchar_t *wstr = PyUnicode_AsWideCharString(unicode, NULL);
+    if (wstr == NULL) {
+        return NULL;
+    }
+
+    size_t error_pos = (size_t)-123;
+    char *str = Py_EncodeLocale(wstr, &error_pos);
+    PyMem_Free(wstr);
+
+    if (str == NULL) {
+        return PyErr_Format(PyExc_ValueError,
+                            "Py_EncodeLocale failed: error_pos=%zd",
+                            error_pos);
+    }
+    assert(error_pos == (size_t)-123);
+
+    PyObject *result = PyBytes_FromString(str);
+    PyMem_Free(str);
+    return result;
+}
+
+
 static PyMethodDef test_methods[] = {
     {"codec_namereplace_errors", codec_namereplace_errors, METH_O},
+    {"decode_locale", decode_locale, METH_O},
+    {"encode_locale", encode_locale, METH_O},
     {NULL},
 };
 
 int
 _PyTestLimitedCAPI_Init_Codec(PyObject *module)
 {
-    if (PyModule_AddFunctions(module, test_methods) < 0) {
-        return -1;
-    }
-    return 0;
+    return PyModule_AddFunctions(module, test_methods);
 }