wtf8

RustPython · youknowone · Jul 15, 2025 · Jul 15, 2025 · Jul 15, 2025 · Jul 15, 2025
commit eedd7d7eddd41be1fd74c3c481ca4e53763a9b07
diff --git a/vm/src/builtins/str.rs b/vm/src/builtins/str.rs
@@ -15,7 +15,7 @@ use crate::{
     format::{format, format_map},
     function::{ArgIterable, ArgSize, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue},
     intern::PyInterned,
-    object::{Traverse, TraverseFn},
+    object::{MaybeTraverse, Traverse, TraverseFn},
     protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
     sequence::SequenceExt,
     sliceable::{SequenceIndex, SliceableSequenceOp},
@@ -350,7 +350,7 @@ impl Constructor for PyStr {
     type Args = StrArgs;
 
     fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
-        let string: PyStrRef = match args.object {
+        let string: PyRef<PyWtf8Str> = match args.object {
             OptionalArg::Present(input) => {
                 if let OptionalArg::Present(enc) = args.encoding {
                     vm.state.codec_registry.decode_text(
@@ -364,7 +364,7 @@ impl Constructor for PyStr {
                 }
             }
             OptionalArg::Missing => {
-                Self::from(String::new()).into_ref_with_type(vm, cls.clone())?
+                Self::from(String::new()).into_ref_with_type(vm, cls.clone())?.into_wtf8()
             }
         };
         if string.class().is(&cls) {
@@ -1499,6 +1499,11 @@ impl PyStrRef {
         self.ensure_valid_utf8(vm)?;
         Ok(unsafe { mem::transmute::<PyRef<PyStr>, PyRef<PyUtf8Str>>(self) })
     }
+
+    pub fn into_wtf8(self) -> PyRef<PyWtf8Str> {
+        // PyStr can always be safely cast to PyWtf8Str
+        unsafe { mem::transmute::<PyRef<PyStr>, PyRef<PyWtf8Str>>(self) }
+    }
 }
 
 impl Representable for PyStr {
@@ -1955,6 +1960,79 @@ impl std::borrow::Borrow<PyObject> for Py<PyUtf8Str> {
     }
 }
 
+#[repr(transparent)]
+#[derive(Debug)]
+pub struct PyWtf8Str(PyStr);
+
+impl PyWtf8Str {
+    //github.com/ Returns the underlying WTF-8 slice.
+    pub fn as_wtf8(&self) -> &Wtf8 {
+        self.0.as_wtf8()
+    }
+}
+
+impl MaybeTraverse for PyWtf8Str {
+    fn try_traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
+        self.0.try_traverse(traverse_fn);
+    }
+}
+
+impl PyPayload for PyWtf8Str {
+    fn class(ctx: &Context) -> &'static Py<PyType> {
+        ctx.types.str_type
+    }
+    fn payload_type_id() -> std::any::TypeId {
+        std::any::TypeId::of::<PyStr>()
+    }
+}
+
+impl From<PyStr> for PyWtf8Str {
+    fn from(s: PyStr) -> Self {
+        PyWtf8Str(s)
+    }
+}
+
+impl<'a> From<&'a str> for PyWtf8Str {
+    fn from(s: &'a str) -> Self {
+        PyWtf8Str(PyStr::from(s))
+    }
+}
+
+impl<'a> From<&'a Wtf8> for PyWtf8Str {
+    fn from(s: &'a Wtf8) -> Self {
+        PyWtf8Str(PyStr::from(s))
+    }
+}
+
+impl From<String> for PyWtf8Str {
+    fn from(s: String) -> Self {
+        PyWtf8Str(PyStr::from(s))
+    }
+}
+
+impl From<Wtf8Buf> for PyWtf8Str {
+    fn from(w: Wtf8Buf) -> Self {
+        PyWtf8Str(PyStr::from(w))
+    }
+}
+
+impl Py<PyWtf8Str> {
+    //github.com/ Upcast to PyStr.
+    pub fn as_pystr(&self) -> &Py<PyStr> {
+        unsafe {
+            // Safety: PyWtf8Str is a wrapper around PyStr, so this cast is safe.
+            &*(self as *const Self as *const Py<PyStr>)
+        }
+    }
+}
+
+impl PartialEq for PyWtf8Str {
+    fn eq(&self, other: &Self) -> bool {
+        self.as_wtf8() == other.as_wtf8()
+    }
+}
+impl Eq for PyWtf8Str {}
+
 impl AnyStrContainer<str> for String {
     fn new() -> Self {
         Self::new()

diff --git a/vm/src/codecs.rs b/vm/src/codecs.rs
@@ -326,7 +326,7 @@ impl CodecsRegistry {
         encoding: &str,
         errors: Option<PyStrRef>,
         vm: &VirtualMachine,
-    ) -> PyResult<PyStrRef> {
+    ) -> PyResult<PyRef<PyWtf8Str>> {
         let codec = self._lookup_text_encoding(encoding, "codecs.decode()", vm)?;
         codec.decode(obj, errors, vm)?.downcast().map_err(|obj| {
             vm.new_type_error(format!(

diff --git a/vm/src/protocol/object.rs b/vm/src/protocol/object.rs
@@ -352,8 +352,8 @@ impl PyObject {
     }
 
     // Container of the virtual machine state:
-    pub fn str(&self, vm: &VirtualMachine) -> PyResult<PyStrRef> {
-        let obj = match self.to_owned().downcast_exact::<PyStr>(vm) {
+    pub fn str(&self, vm: &VirtualMachine) -> PyResult<PyRef<PyWtf8Str>> {
+        let obj = match self.to_owned().downcast_exact::<PyWtf8Str>(vm) {
             Ok(s) => return Ok(s.into_pyref()),
             Err(obj) => obj,
         };
@@ -363,7 +363,7 @@ impl PyObject {
             None => return obj.repr(vm),
         };
         let s = str_method.invoke((), vm)?;
-        s.downcast::<PyStr>().map_err(|obj| {
+        s.downcast::<PyWtf8Str>().map_err(|obj| {
             vm.new_type_error(format!(
                 "__str__ returned non-string (type {})",
                 obj.class().name()