Make cformat wtf8-compatible

Lib/test/test_codecs.py

@@ -1698,8 +1698,6 @@ nameprep_tests = [
 
 
 class NameprepTest(unittest.TestCase):
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_nameprep(self):
         from encodings.idna import nameprep
         for pos, (orig, prepped) in enumerate(nameprep_tests):

common/src/cformat.rs

@@ -11,11 +11,13 @@ use std::{
     str::FromStr,
 };
 
+use crate::wtf8::{CodePoint, Wtf8, Wtf8Buf};
+
 #[derive(Debug, PartialEq)]
 pub enum CFormatErrorType {
     UnmatchedKeyParentheses,
     MissingModuloSign,
-    UnsupportedFormatChar(char),
+    UnsupportedFormatChar(CodePoint),
     IncompleteFormat,
     IntTooBig,
     // Unimplemented,
@@ -39,7 +41,9 @@ impl fmt::Display for CFormatError {
             UnsupportedFormatChar(c) => write!(
                 f,
                 "unsupported format character '{}' ({:#x}) at index {}",
-                c, c as u32, self.index
+                c,
+                c.to_u32(),
+                self.index
             ),
             IntTooBig => write!(f, "width/precision too big"),
             _ => write!(f, "unexpected error parsing format string"),
@@ -160,7 +164,7 @@ pub trait FormatBuf:
     fn concat(self, other: Self) -> Self;
 }
 
-pub trait FormatChar: Copy + Into<char> + From<u8> {
+pub trait FormatChar: Copy + Into<CodePoint> + From<u8> {
     fn to_char_lossy(self) -> char;
     fn eq_char(self, c: char) -> bool;
 }
@@ -188,6 +192,29 @@ impl FormatChar for char {
     }
 }
 
+impl FormatBuf for Wtf8Buf {
+    type Char = CodePoint;
+    fn chars(&self) -> impl Iterator<Item = Self::Char> {
+        self.code_points()
+    }
+    fn len(&self) -> usize {
+        (**self).len()
+    }
+    fn concat(mut self, other: Self) -> Self {
+        self.extend([other]);
+        self
+    }
+}
+
+impl FormatChar for CodePoint {
+    fn to_char_lossy(self) -> char {
+        self.to_char_lossy()
+    }
+    fn eq_char(self, c: char) -> bool {
+        self == c
+    }
+}
+
 impl FormatBuf for Vec<u8> {
     type Char = u8;
     fn chars(&self) -> impl Iterator<Item = Self::Char> {
@@ -801,6 +828,15 @@ impl FromStr for CFormatString {
     }
 }
 
+pub type CFormatWtf8 = CFormatStrOrBytes<Wtf8Buf>;
+
+impl CFormatWtf8 {
+    pub fn parse_from_wtf8(s: &Wtf8) -> Result<Self, CFormatError> {
+        let mut iter = s.code_points().enumerate().peekable();
+        Self::parse(&mut iter)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

common/src/wtf8/mod.rs

@@ -75,6 +75,12 @@ impl fmt::Debug for CodePoint {
     }
 }
 
+impl fmt::Display for CodePoint {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.to_char_lossy().fmt(f)
+    }
+}
+
 impl CodePoint {
     /// Unsafely creates a new `CodePoint` without checking the value.
     ///
@@ -109,13 +115,13 @@ impl CodePoint {
 
     /// Returns the numeric value of the code point.
     #[inline]
-    pub fn to_u32(&self) -> u32 {
+    pub fn to_u32(self) -> u32 {
         self.value
     }
 
     /// Returns the numeric value of the code point if it is a leading surrogate.
     #[inline]
-    pub fn to_lead_surrogate(&self) -> Option<u16> {
+    pub fn to_lead_surrogate(self) -> Option<u16> {
         match self.value {
             lead @ 0xD800..=0xDBFF => Some(lead as u16),
             _ => None,
@@ -124,7 +130,7 @@ impl CodePoint {
 
     /// Returns the numeric value of the code point if it is a trailing surrogate.
     #[inline]
-    pub fn to_trail_surrogate(&self) -> Option<u16> {
+    pub fn to_trail_surrogate(self) -> Option<u16> {
         match self.value {
             trail @ 0xDC00..=0xDFFF => Some(trail as u16),
             _ => None,
@@ -135,7 +141,7 @@ impl CodePoint {
     ///
     /// Returns `None` if the code point is a surrogate (from U+D800 to U+DFFF).
     #[inline]
-    pub fn to_char(&self) -> Option<char> {
+    pub fn to_char(self) -> Option<char> {
         match self.value {
             0xD800..=0xDFFF => None,
             _ => Some(unsafe { char::from_u32_unchecked(self.value) }),
@@ -147,7 +153,7 @@ impl CodePoint {
     /// Returns `'\u{FFFD}'` (the replacement character “�”)
     /// if the code point is a surrogate (from U+D800 to U+DFFF).
     #[inline]
-    pub fn to_char_lossy(&self) -> char {
+    pub fn to_char_lossy(self) -> char {
         self.to_char().unwrap_or('\u{FFFD}')
     }
 
@@ -170,6 +176,12 @@ impl From<u16> for CodePoint {
     }
 }
 
+impl From<u8> for CodePoint {
+    fn from(value: u8) -> Self {
+        char::from(value).into()
+    }
+}
+
 impl From<char> for CodePoint {
     fn from(value: char) -> Self {
         Self::from_char(value)
@@ -515,6 +527,13 @@ impl Extend<CodePoint> for Wtf8Buf {
     }
 }
 
+impl<W: AsRef<Wtf8>> Extend<W> for Wtf8Buf {
+    fn extend<T: IntoIterator<Item = W>>(&mut self, iter: T) {
+        iter.into_iter()
+            .for_each(move |w| self.push_wtf8(w.as_ref()));
+    }
+}
+
 impl<W: AsRef<Wtf8>> FromIterator<W> for Wtf8Buf {
     fn from_iter<T: IntoIterator<Item = W>>(iter: T) -> Self {
         let mut buf = Wtf8Buf::new();
@@ -541,6 +560,12 @@ impl From<&str> for Wtf8Buf {
     }
 }
 
+impl From<ascii::AsciiString> for Wtf8Buf {
+    fn from(s: ascii::AsciiString) -> Self {
+        Wtf8Buf::from_string(s.into())
+    }
+}
+
 /// A borrowed slice of well-formed WTF-8 data.
 ///
 /// Similar to `&str`, but can additionally contain surrogate code points

vm/src/builtins/str.rs

@@ -857,8 +857,8 @@ impl PyStr {
     }
 
     #[pymethod(name = "__mod__")]
-    fn modulo(&self, values: PyObjectRef, vm: &VirtualMachine) -> PyResult<String> {
-        cformat_string(vm, self.as_str(), values)
+    fn modulo(&self, values: PyObjectRef, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
+        cformat_string(vm, self.as_wtf8(), values)
     }
 
     #[pymethod(magic)]

vm/src/cformat.rs

@@ -2,6 +2,7 @@
 //! as per the [Python Docs](https://docs.python.org/3/library/stdtypes.html#printf-style-string-formatting).
 
 use crate::common::cformat::*;
+use crate::common::wtf8::{CodePoint, Wtf8, Wtf8Buf};
 use crate::{
     AsObject, PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject, VirtualMachine,
     builtins::{
@@ -125,13 +126,13 @@ fn spec_format_string(
     spec: &CFormatSpec,
     obj: PyObjectRef,
     idx: usize,
-) -> PyResult<String> {
+) -> PyResult<Wtf8Buf> {
     match &spec.format_type {
         CFormatType::String(conversion) => {
             let result = match conversion {
                 CFormatConversion::Ascii => builtins::ascii(obj, vm)?.into(),
-                CFormatConversion::Str => obj.str(vm)?.as_str().to_owned(),
-                CFormatConversion::Repr => obj.repr(vm)?.as_str().to_owned(),
+                CFormatConversion::Str => obj.str(vm)?.as_wtf8().to_owned(),
+                CFormatConversion::Repr => obj.repr(vm)?.as_wtf8().to_owned(),
                 CFormatConversion::Bytes => {
                     // idx is the position of the %, we want the position of the b
                     return Err(vm.new_value_error(format!(
@@ -146,16 +147,18 @@ fn spec_format_string(
             CNumberType::DecimalD | CNumberType::DecimalI | CNumberType::DecimalU => {
                 match_class!(match &obj {
                     ref i @ PyInt => {
-                        Ok(spec.format_number(i.as_bigint()))
+                        Ok(spec.format_number(i.as_bigint()).into())
                     }
                     ref f @ PyFloat => {
-                        Ok(spec.format_number(&try_f64_to_bigint(f.to_f64(), vm)?))
+                        Ok(spec
+                            .format_number(&try_f64_to_bigint(f.to_f64(), vm)?)
+                            .into())
                     }
                     obj => {
                         if let Some(method) = vm.get_method(obj.clone(), identifier!(vm, __int__)) {
                             let result = method?.call((), vm)?;
                             if let Some(i) = result.payload::<PyInt>() {
-                                return Ok(spec.format_number(i.as_bigint()));
+                                return Ok(spec.format_number(i.as_bigint()).into());
                             }
                         }
                         Err(vm.new_type_error(format!(
@@ -168,7 +171,7 @@ fn spec_format_string(
             }
             _ => {
                 if let Some(i) = obj.payload::<PyInt>() {
-                    Ok(spec.format_number(i.as_bigint()))
+                    Ok(spec.format_number(i.as_bigint()).into())
                 } else {
                     Err(vm.new_type_error(format!(
                         "%{} format: an integer is required, not {}",
@@ -180,21 +183,21 @@ fn spec_format_string(
         },
         CFormatType::Float(_) => {
             let value = ArgIntoFloat::try_from_object(vm, obj)?;
-            Ok(spec.format_float(value.into()))
+            Ok(spec.format_float(value.into()).into())
         }
         CFormatType::Character(CCharacterType::Character) => {
             if let Some(i) = obj.payload::<PyInt>() {
                 let ch = i
                     .as_bigint()
                     .to_u32()
-                    .and_then(char::from_u32)
+                    .and_then(CodePoint::from_u32)
                     .ok_or_else(|| {
                         vm.new_overflow_error("%c arg not in range(0x110000)".to_owned())
                     })?;
                 return Ok(spec.format_char(ch));
             }
             if let Some(s) = obj.payload::<PyStr>() {
-                if let Ok(ch) = s.as_str().chars().exactly_one() {
+                if let Ok(ch) = s.as_wtf8().code_points().exactly_one() {
                     return Ok(spec.format_char(ch));
                 }
             }
@@ -374,17 +377,16 @@ pub(crate) fn cformat_bytes(
 
 pub(crate) fn cformat_string(
     vm: &VirtualMachine,
-    format_string: &str,
+    format_string: &Wtf8,
     values_obj: PyObjectRef,
-) -> PyResult<String> {
-    let format = format_string
-        .parse::<CFormatString>()
+) -> PyResult<Wtf8Buf> {
+    let format = CFormatWtf8::parse_from_wtf8(format_string)
         .map_err(|err| vm.new_value_error(err.to_string()))?;
     let (num_specifiers, mapping_required) = format
         .check_specifiers()
         .ok_or_else(|| specifier_error(vm))?;
 
-    let mut result = String::new();
+    let mut result = Wtf8Buf::new();
 
     let is_mapping = values_obj.class().has_attr(identifier!(vm, __getitem__))
         && !values_obj.fast_isinstance(vm.ctx.types.tuple_type)
@@ -399,7 +401,7 @@ pub(crate) fn cformat_string(
         {
             for (_, part) in format.iter() {
                 match part {
-                    CFormatPart::Literal(literal) => result.push_str(literal),
+                    CFormatPart::Literal(literal) => result.push_wtf8(literal),
                     CFormatPart::Spec(_) => unreachable!(),
                 }
             }
@@ -415,11 +417,11 @@ pub(crate) fn cformat_string(
         return if is_mapping {
             for (idx, part) in format {
                 match part {
-                    CFormatPart::Literal(literal) => result.push_str(&literal),
+                    CFormatPart::Literal(literal) => result.push_wtf8(&literal),
                     CFormatPart::Spec(CFormatSpecKeyed { mapping_key, spec }) => {
                         let value = values_obj.get_item(&mapping_key.unwrap(), vm)?;
                         let part_result = spec_format_string(vm, &spec, value, idx)?;
-                        result.push_str(&part_result);
+                        result.push_wtf8(&part_result);
                     }
                 }
             }
@@ -439,7 +441,7 @@ pub(crate) fn cformat_string(
 
     for (idx, part) in format {
         match part {
-            CFormatPart::Literal(literal) => result.push_str(&literal),
+            CFormatPart::Literal(literal) => result.push_wtf8(&literal),
             CFormatPart::Spec(CFormatSpecKeyed { mut spec, .. }) => {
                 try_update_quantity_from_tuple(
                     vm,
@@ -456,7 +458,7 @@ pub(crate) fn cformat_string(
                     }
                 }?;
                 let part_result = spec_format_string(vm, &spec, value, idx)?;
-                result.push_str(&part_result);
+                result.push_wtf8(&part_result);
             }
         }
     }

vm/src/dictdatatype.rs

@@ -12,6 +12,7 @@ use crate::{
     common::{
         hash,
         lock::{PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard},
+        wtf8::{Wtf8, Wtf8Buf},
     },
     object::{Traverse, TraverseFn},
 };
@@ -750,7 +751,7 @@ impl DictKey for Py<PyStr> {
         if self.is(other_key) {
             Ok(true)
         } else if let Some(pystr) = other_key.payload_if_exact::<PyStr>(vm) {
-            Ok(pystr.as_str() == self.as_str())
+            Ok(self.as_wtf8() == pystr.as_wtf8())
         } else {
             vm.bool_eq(self.as_object(), other_key)
         }
@@ -834,7 +835,7 @@ impl DictKey for str {
 
     fn key_eq(&self, vm: &VirtualMachine, other_key: &PyObject) -> PyResult<bool> {
         if let Some(pystr) = other_key.payload_if_exact::<PyStr>(vm) {
-            Ok(pystr.as_str() == self)
+            Ok(pystr.as_wtf8() == self.as_ref())
         } else {
             // Fall back to PyObjectRef implementation.
             let s = vm.ctx.new_str(self);
@@ -871,6 +872,63 @@ impl DictKey for String {
     }
 }
 
+impl DictKey for Wtf8 {
+    type Owned = Wtf8Buf;
+    #[inline(always)]
+    fn _to_owned(&self, _vm: &VirtualMachine) -> Self::Owned {
+        self.to_owned()
+    }
+    #[inline]
+    fn key_hash(&self, vm: &VirtualMachine) -> PyResult<HashValue> {
+        // follow a similar route as the hashing of PyStrRef
+        Ok(vm.state.hash_secret.hash_bytes(self.as_bytes()))
+    }
+    #[inline(always)]
+    fn key_is(&self, _other: &PyObject) -> bool {
+        // No matter who the other pyobject is, we are never the same thing, since
+        // we are a str, not a pyobject.
+        false
+    }
+
+    fn key_eq(&self, vm: &VirtualMachine, other_key: &PyObject) -> PyResult<bool> {
+        if let Some(pystr) = other_key.payload_if_exact::<PyStr>(vm) {
+            Ok(pystr.as_wtf8() == self)
+        } else {
+            // Fall back to PyObjectRef implementation.
+            let s = vm.ctx.new_str(self);
+            s.key_eq(vm, other_key)
+        }
+    }
+
+    fn key_as_isize(&self, vm: &VirtualMachine) -> PyResult<isize> {
+        Err(vm.new_type_error("'str' object cannot be interpreted as an integer".to_owned()))
+    }
+}
+
+impl DictKey for Wtf8Buf {
+    type Owned = Wtf8Buf;
+    #[inline]
+    fn _to_owned(&self, _vm: &VirtualMachine) -> Self::Owned {
+        self.clone()
+    }
+
+    fn key_hash(&self, vm: &VirtualMachine) -> PyResult<HashValue> {
+        (**self).key_hash(vm)
+    }
+
+    fn key_is(&self, other: &PyObject) -> bool {
+        (**self).key_is(other)
+    }
+
+    fn key_eq(&self, vm: &VirtualMachine, other_key: &PyObject) -> PyResult<bool> {
+        (**self).key_eq(vm, other_key)
+    }
+
+    fn key_as_isize(&self, vm: &VirtualMachine) -> PyResult<isize> {
+        (**self).key_as_isize(vm)
+    }
+}
+
 impl DictKey for [u8] {
     type Owned = Vec<u8>;
     #[inline(always)]