Accept surrogates in _json.encode_basestring{,_ascii} (#7673)

encode_basestring/encode_basestring_ascii took PyUtf8StrRef, so
json.dumps(str_with_lone_surrogate) raised UnicodeEncodeError at the
Python/Rust boundary before write_json_string ran. CPython's encoder
emits \uXXXX under ensure_ascii=True and passes raw WTF-8 otherwise.

Switch to PyStrRef + s.as_wtf8(), matching scanstring in the same file.
Rewrite write_json_string to accept &Wtf8 and iterate
code_point_indices, emitting \uXXXX for surrogates in ascii mode and
passing their bytes through otherwise. Stop escaping 0x7F in the
ensure_ascii=False path (matches py_encode_basestring). Return Wtf8Buf
via the checked from_bytes so invariant breaks panic instead of UB.

Fuzzing also exposed two pre-existing ESCAPE_CHARS typos: 0x0B was
"\u000" and 0x1B was "\u001" (both missing trailing 'b'). Fixed here.

Verified byte-identical with CPython 3.13.4 over 16 manual + 10,000
random fuzz cases. Full test.test_json: 214 tests, 0 failures, 0
unexpected successes. Unmasks test_ascii_non_printable_encode and
test_single_surrogate_encode. Decoder path is a follow-up.

Lib/test/test_json/test_unicode.py

@@ -138,14 +138,6 @@ class TestUnicode:
 
 class TestPyUnicode(TestUnicode, PyTest): pass
 class TestCUnicode(TestUnicode, CTest):
-    @unittest.expectedFailure  # TODO: RUSTPYTHON
-    def test_ascii_non_printable_encode(self):
-        return super().test_ascii_non_printable_encode()
-
-    @unittest.skip("TODO: RUSTPYTHON; panics with 'str has surrogates'")
+    @unittest.skip("TODO: RUSTPYTHON; decode path still uses PyUtf8StrRef")
     def test_single_surrogate_decode(self):
         return super().test_single_surrogate_decode()
-
-    @unittest.skip("TODO: RUSTPYTHON; panics with 'str has surrogates'")
-    def test_single_surrogate_encode(self):
-        return super().test_single_surrogate_encode()

crates/stdlib/src/json.rs

@@ -667,24 +667,30 @@ mod _json {
         }
     }
 
-    fn encode_string(s: &str, ascii_only: bool) -> String {
+    fn encode_string(wtf8: &rustpython_common::wtf8::Wtf8, ascii_only: bool) -> Wtf8Buf {
         flame_guard!("_json::encode_string");
-        let mut buf = Vec::<u8>::with_capacity(s.len() + 2);
-        machinery::write_json_string(s, ascii_only, &mut buf)
+        let mut buf = Vec::<u8>::with_capacity(wtf8.len() + 2);
+        machinery::write_json_string(wtf8, ascii_only, &mut buf)
             // SAFETY: writing to a vec can't fail
             .unwrap_or_else(|_| unsafe { core::hint::unreachable_unchecked() });
-        // SAFETY: we only output valid utf8 from write_json_string
-        unsafe { String::from_utf8_unchecked(buf) }
+        // write_json_string is designed to produce valid WTF-8 bytes:
+        // - ASCII control characters and JSON-specials are written as ASCII escapes
+        // - Valid Unicode scalars are written as UTF-8 (a subset of WTF-8)
+        // - Lone surrogates (ascii_only=false branch only) pass through as the
+        //   input's WTF-8 byte sequences unchanged
+        // Use the checked constructor so any violation of that invariant
+        // surfaces as a panic during testing instead of undefined behavior.
+        Wtf8Buf::from_bytes(buf).expect("write_json_string produced invalid WTF-8")
     }
 
     #[pyfunction]
-    fn encode_basestring(s: PyUtf8StrRef) -> String {
-        encode_string(s.as_str(), false)
+    fn encode_basestring(s: PyStrRef) -> Wtf8Buf {
+        encode_string(s.as_wtf8(), false)
     }
 
     #[pyfunction]
-    fn encode_basestring_ascii(s: PyUtf8StrRef) -> String {
-        encode_string(s.as_str(), true)
+    fn encode_basestring_ascii(s: PyStrRef) -> Wtf8Buf {
+        encode_string(s.as_wtf8(), true)
     }
 
     fn py_decode_error(

crates/stdlib/src/json/machinery.rs

@@ -35,9 +35,9 @@ use rustpython_common::wtf8::{CodePoint, Wtf8, Wtf8Buf};
 
 static ESCAPE_CHARS: [&str; 0x20] = [
     "\\u0000", "\\u0001", "\\u0002", "\\u0003", "\\u0004", "\\u0005", "\\u0006", "\\u0007", "\\b",
-    "\\t", "\\n", "\\u000", "\\f", "\\r", "\\u000e", "\\u000f", "\\u0010", "\\u0011", "\\u0012",
+    "\\t", "\\n", "\\u000b", "\\f", "\\r", "\\u000e", "\\u000f", "\\u0010", "\\u0011", "\\u0012",
     "\\u0013", "\\u0014", "\\u0015", "\\u0016", "\\u0017", "\\u0018", "\\u0019", "\\u001a",
-    "\\u001", "\\u001c", "\\u001d", "\\u001e", "\\u001f",
+    "\\u001b", "\\u001c", "\\u001d", "\\u001e", "\\u001f",
 ];
 
 // This bitset represents which bytes can be copied as-is to a JSON string (0)
@@ -72,30 +72,51 @@ fn json_escaped_char(c: u8) -> Option<&'static str> {
     }
 }
 
-pub fn write_json_string<W: io::Write>(s: &str, ascii_only: bool, w: &mut W) -> io::Result<()> {
+pub fn write_json_string<W: io::Write>(wtf8: &Wtf8, ascii_only: bool, w: &mut W) -> io::Result<()> {
     w.write_all(b"\"")?;
     let mut write_start_idx = 0;
-    let bytes = s.as_bytes();
+    let bytes = wtf8.as_bytes();
     if ascii_only {
-        for (idx, c) in s.char_indices() {
-            if c.is_ascii() {
-                if let Some(escaped) = json_escaped_char(c as u8) {
+        for (idx, cp) in wtf8.code_point_indices() {
+            if let Some(c) = cp.to_char() {
+                // Valid Unicode scalar.
+                if c.is_ascii() {
+                    if let Some(escaped) = json_escaped_char(c as u8) {
+                        w.write_all(&bytes[write_start_idx..idx])?;
+                        w.write_all(escaped.as_bytes())?;
+                        write_start_idx = idx + 1;
+                    }
+                } else {
                     w.write_all(&bytes[write_start_idx..idx])?;
-                    w.write_all(escaped.as_bytes())?;
-                    write_start_idx = idx + 1;
+                    write_start_idx = idx + c.len_utf8();
+                    // codepoints outside the BMP get 2 '\uxxxx' sequences to represent them
+                    for point in c.encode_utf16(&mut [0; 2]) {
+                        write!(w, "\\u{point:04x}")?;
+                    }
                 }
             } else {
+                // Lone surrogate code point (U+D800..U+DFFF).
+                // WTF-8 encodes these as 3-byte sequences; skip those raw bytes
+                // and emit a \uXXXX escape with the surrogate value.
                 w.write_all(&bytes[write_start_idx..idx])?;
-                write_start_idx = idx + c.len_utf8();
-                // codepoints outside the BMP get 2 '\uxxxx' sequences to represent them
-                for point in c.encode_utf16(&mut [0; 2]) {
-                    write!(w, "\\u{point:04x}")?;
-                }
+                write_start_idx = idx + 3;
+                write!(w, "\\u{:04x}", cp.to_u32())?;
             }
         }
     } else {
-        for (idx, c) in s.bytes().enumerate() {
-            if let Some(escaped) = json_escaped_char(c) {
+        // ensure_ascii is false: only JSON-required escapes (< 0x20, \, ")
+        // are applied. 0x7F (DEL) is NOT escaped here, matching CPython's
+        // py_encode_basestring. Multi-byte UTF-8 characters and WTF-8
+        // surrogate sequences flow through unchanged via the trailing flush,
+        // so surrogates round-trip as-is (matching CPython behavior).
+        for (idx, c) in wtf8.as_bytes().iter().enumerate() {
+            let escaped_opt: Option<&'static str> = match *c {
+                x if x < 0x20 => Some(ESCAPE_CHARS[x as usize]),
+                b'\\' => Some("\\\\"),
+                b'\"' => Some("\\\""),
+                _ => None,
+            };
+            if let Some(escaped) = escaped_opt {
                 w.write_all(&bytes[write_start_idx..idx])?;
                 w.write_all(escaped.as_bytes())?;
                 write_start_idx = idx + 1;