Implement _json.encode_basestring{,_ascii}

coolreader18 · coolreader18 · commit 0fb79e1086ba · 2020-06-06T15:33:29.000-05:00
diff --git a/Lib/test/test_json/test_speedups.py b/Lib/test/test_json/test_speedups.py
@@ -14,8 +14,6 @@ def test_scanstring(self):
         self.assertEqual(self.json.decoder.scanstring.__module__, "_json")
         self.assertIs(self.json.decoder.scanstring, self.json.decoder.c_scanstring)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_encode_basestring_ascii(self):
         self.assertEqual(self.json.encoder.encode_basestring_ascii.__module__,
                          "_json")
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -120,7 +120,7 @@ libz-sys = "1.0"
 winreg = "0.7"
 schannel = "0.1"
 
-[target."cfg(windows)".dependencies.winapi]
+[target.'cfg(windows)'.dependencies.winapi]
 version = "0.3"
 features = ["winsock2", "handleapi", "ws2def", "std", "winbase", "wincrypt", "fileapi"]
 
diff --git a/vm/src/stdlib/json.rs b/vm/src/stdlib/json.rs
@@ -7,6 +7,8 @@ use crate::VirtualMachine;
 use num_bigint::BigInt;
 use std::str::FromStr;
 
+mod machinery;
+
 #[pyclass(name = "Scanner")]
 #[derive(Debug)]
 struct JsonScanner {
@@ -209,11 +211,30 @@ impl JsonScanner {
     }
 }
 
+fn encode_string(s: &str, ascii_only: bool) -> String {
+    let mut buf = Vec::<u8>::with_capacity(s.len() + 2);
+    machinery::write_json_string(s, ascii_only, &mut buf)
+        // writing to a vec can't fail
+        .unwrap_or_else(|_| unsafe { std::hint::unreachable_unchecked() });
+    // TODO: verify that the implementation is correct enough to use `from_utf8_unchecked`
+    String::from_utf8(buf).expect("invalid utf-8 in json output")
+}
+
+fn _json_encode_basestring(s: PyStringRef) -> String {
+    encode_string(s.as_str(), false)
+}
+
+fn _json_encode_basestring_ascii(s: PyStringRef) -> String {
+    encode_string(s.as_str(), true)
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
     let scanner_cls = JsonScanner::make_class(ctx);
     scanner_cls.set_str_attr("__module__", vm.new_str("_json".to_owned()));
     py_module!(vm, "_json", {
         "make_scanner" => scanner_cls,
+        "encode_basestring" => named_function!(ctx, _json, encode_basestring),
+        "encode_basestring_ascii" => named_function!(ctx, _json, encode_basestring_ascii),
     })
 }
diff --git a/vm/src/stdlib/json/machinery.rs b/vm/src/stdlib/json/machinery.rs
@@ -0,0 +1,98 @@
+// derived from https://github.com/lovasoa/json_in_type
+
+// BSD 2-Clause License
+//
+// Copyright (c) 2018, Ophir LOJKINE
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice, this
+//   list of conditions and the following disclaimer.
+//
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+use std::io;
+
+static ESCAPE_CHARS: [&[u8]; 0x20] = [
+    b"\\u0000", b"\\u0001", b"\\u0002", b"\\u0003", b"\\u0004", b"\\u0005", b"\\u0006", b"\\u0007",
+    b"\\b", b"\\t", b"\\n", b"\\u000b", b"\\f", b"\\r", b"\\u000e", b"\\u000f", b"\\u0010",
+    b"\\u0011", b"\\u0012", b"\\u0013", b"\\u0014", b"\\u0015", b"\\u0016", b"\\u0017", b"\\u0018",
+    b"\\u0019", b"\\u001a", b"\\u001b", b"\\u001c", b"\\u001d", b"\\u001e", b"\\u001f",
+];
+
+// This bitset represents which bytes can be copied as-is to a JSON string (0)
+// And which one need to be escaped (1)
+// The characters that need escaping are 0x00 to 0x1F, 0x22 ("), 0x5C (\), 0x7F (DEL)
+// Non-ASCII unicode characters can be safely included in a JSON string
+static NEEDS_ESCAPING_BITSET: [u64; 4] = [
+    //fedcba9876543210_fedcba9876543210_fedcba9876543210_fedcba9876543210
+    0b0000000000000000_0000000000000100_1111111111111111_1111111111111111, // 3_2_1_0
+    0b1000000000000000_0000000000000000_0001000000000000_0000000000000000, // 7_6_5_4
+    0b0000000000000000_0000000000000000_0000000000000000_0000000000000000, // B_A_9_8
+    0b0000000000000000_0000000000000000_0000000000000000_0000000000000000, // F_E_D_C
+];
+
+#[inline(always)]
+fn json_escaped_char(c: u8) -> Option<&'static [u8]> {
+    let bitset_value = NEEDS_ESCAPING_BITSET[(c / 64) as usize] & (1 << (c % 64));
+    if bitset_value == 0 {
+        None
+    } else {
+        Some(match c {
+            x if x < 0x20 => ESCAPE_CHARS[c as usize],
+            b'\\' => &b"\\\\"[..],
+            b'\"' => &b"\\\""[..],
+            0x7F => &b"\\u007f"[..],
+            _ => unreachable!(),
+        })
+    }
+}
+
+pub fn write_json_string<W: io::Write>(s: &str, ascii_only: bool, w: &mut W) -> io::Result<()> {
+    w.write_all(b"\"")?;
+    let mut write_start_idx = 0;
+    let bytes = s.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) {
+                    w.write_all(&bytes[write_start_idx..idx])?;
+                    w.write_all(escaped)?;
+                    write_start_idx = idx + 1;
+                }
+            } else {
+                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{:04x}", point)?;
+                }
+            }
+        }
+    } else {
+        for (idx, c) in s.bytes().enumerate() {
+            if let Some(escaped) = json_escaped_char(c) {
+                w.write_all(&bytes[write_start_idx..idx])?;
+                w.write_all(escaped)?;
+                write_start_idx = idx + 1;
+            }
+        }
+    }
+    w.write_all(&bytes[write_start_idx..])?;
+    w.write_all(b"\"")
+}
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -1483,13 +1483,6 @@ impl VirtualMachine {
         attr_name: impl TryIntoRef<PyString>,
         attr_value: impl Into<PyObjectRef>,
     ) -> PyResult<()> {
-        // let attr_name = attr_name.try_into_ref(self)?;
-        // let value = attr_value.into();
-        // let dict = module.dict().expect("module doesn't have dict");
-        // if let Ok(module_name) = dict.get_item("__name__", self) {
-        //     let _ = self.set_attr(&value, "__module__", module_name);
-        // }
-        // dict.set_item(&attr_name, value, self)?;
         let val = attr_value.into();
         objobject::setattr(module.clone(), attr_name.try_into_ref(self)?, val, self)
     }
