Skip to content

Commit eedd7d7

Browse files
committed
wtf8
1 parent 6946a14 commit eedd7d7

3 files changed

Lines changed: 85 additions & 7 deletions

File tree

‎vm/src/builtins/str.rs‎

Lines changed: 81 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ use crate::{
1515
format::{format, format_map},
1616
function::{ArgIterable, ArgSize, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue},
1717
intern::PyInterned,
18-
object::{Traverse, TraverseFn},
18+
object::{MaybeTraverse, Traverse, TraverseFn},
1919
protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
2020
sequence::SequenceExt,
2121
sliceable::{SequenceIndex, SliceableSequenceOp},
@@ -350,7 +350,7 @@ impl Constructor for PyStr {
350350
type Args = StrArgs;
351351

352352
fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
353-
let string: PyStrRef = match args.object {
353+
let string: PyRef<PyWtf8Str> = match args.object {
354354
OptionalArg::Present(input) => {
355355
if let OptionalArg::Present(enc) = args.encoding {
356356
vm.state.codec_registry.decode_text(
@@ -364,7 +364,7 @@ impl Constructor for PyStr {
364364
}
365365
}
366366
OptionalArg::Missing => {
367-
Self::from(String::new()).into_ref_with_type(vm, cls.clone())?
367+
Self::from(String::new()).into_ref_with_type(vm, cls.clone())?.into_wtf8()
368368
}
369369
};
370370
if string.class().is(&cls) {
@@ -1499,6 +1499,11 @@ impl PyStrRef {
14991499
self.ensure_valid_utf8(vm)?;
15001500
Ok(unsafe { mem::transmute::<PyRef<PyStr>, PyRef<PyUtf8Str>>(self) })
15011501
}
1502+
1503+
pub fn into_wtf8(self) -> PyRef<PyWtf8Str> {
1504+
// PyStr can always be safely cast to PyWtf8Str
1505+
unsafe { mem::transmute::<PyRef<PyStr>, PyRef<PyWtf8Str>>(self) }
1506+
}
15021507
}
15031508

15041509
impl Representable for PyStr {
@@ -1955,6 +1960,79 @@ impl std::borrow::Borrow<PyObject> for Py<PyUtf8Str> {
19551960
}
19561961
}
19571962

1963+
#[repr(transparent)]
1964+
#[derive(Debug)]
1965+
pub struct PyWtf8Str(PyStr);
1966+
1967+
impl PyWtf8Str {
1968+
/// Returns the underlying WTF-8 slice.
1969+
pub fn as_wtf8(&self) -> &Wtf8 {
1970+
self.0.as_wtf8()
1971+
}
1972+
}
1973+
1974+
impl MaybeTraverse for PyWtf8Str {
1975+
fn try_traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
1976+
self.0.try_traverse(traverse_fn);
1977+
}
1978+
}
1979+
1980+
impl PyPayload for PyWtf8Str {
1981+
fn class(ctx: &Context) -> &'static Py<PyType> {
1982+
ctx.types.str_type
1983+
}
1984+
fn payload_type_id() -> std::any::TypeId {
1985+
std::any::TypeId::of::<PyStr>()
1986+
}
1987+
}
1988+
1989+
impl From<PyStr> for PyWtf8Str {
1990+
fn from(s: PyStr) -> Self {
1991+
PyWtf8Str(s)
1992+
}
1993+
}
1994+
1995+
impl<'a> From<&'a str> for PyWtf8Str {
1996+
fn from(s: &'a str) -> Self {
1997+
PyWtf8Str(PyStr::from(s))
1998+
}
1999+
}
2000+
2001+
impl<'a> From<&'a Wtf8> for PyWtf8Str {
2002+
fn from(s: &'a Wtf8) -> Self {
2003+
PyWtf8Str(PyStr::from(s))
2004+
}
2005+
}
2006+
2007+
impl From<String> for PyWtf8Str {
2008+
fn from(s: String) -> Self {
2009+
PyWtf8Str(PyStr::from(s))
2010+
}
2011+
}
2012+
2013+
impl From<Wtf8Buf> for PyWtf8Str {
2014+
fn from(w: Wtf8Buf) -> Self {
2015+
PyWtf8Str(PyStr::from(w))
2016+
}
2017+
}
2018+
2019+
impl Py<PyWtf8Str> {
2020+
/// Upcast to PyStr.
2021+
pub fn as_pystr(&self) -> &Py<PyStr> {
2022+
unsafe {
2023+
// Safety: PyWtf8Str is a wrapper around PyStr, so this cast is safe.
2024+
&*(self as *const Self as *const Py<PyStr>)
2025+
}
2026+
}
2027+
}
2028+
2029+
impl PartialEq for PyWtf8Str {
2030+
fn eq(&self, other: &Self) -> bool {
2031+
self.as_wtf8() == other.as_wtf8()
2032+
}
2033+
}
2034+
impl Eq for PyWtf8Str {}
2035+
19582036
impl AnyStrContainer<str> for String {
19592037
fn new() -> Self {
19602038
Self::new()

‎vm/src/codecs.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -326,7 +326,7 @@ impl CodecsRegistry {
326326
encoding: &str,
327327
errors: Option<PyStrRef>,
328328
vm: &VirtualMachine,
329-
) -> PyResult<PyStrRef> {
329+
) -> PyResult<PyRef<PyWtf8Str>> {
330330
let codec = self._lookup_text_encoding(encoding, "codecs.decode()", vm)?;
331331
codec.decode(obj, errors, vm)?.downcast().map_err(|obj| {
332332
vm.new_type_error(format!(

‎vm/src/protocol/object.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -352,8 +352,8 @@ impl PyObject {
352352
}
353353

354354
// Container of the virtual machine state:
355-
pub fn str(&self, vm: &VirtualMachine) -> PyResult<PyStrRef> {
356-
let obj = match self.to_owned().downcast_exact::<PyStr>(vm) {
355+
pub fn str(&self, vm: &VirtualMachine) -> PyResult<PyRef<PyWtf8Str>> {
356+
let obj = match self.to_owned().downcast_exact::<PyWtf8Str>(vm) {
357357
Ok(s) => return Ok(s.into_pyref()),
358358
Err(obj) => obj,
359359
};
@@ -363,7 +363,7 @@ impl PyObject {
363363
None => return obj.repr(vm),
364364
};
365365
let s = str_method.invoke((), vm)?;
366-
s.downcast::<PyStr>().map_err(|obj| {
366+
s.downcast::<PyWtf8Str>().map_err(|obj| {
367367
vm.new_type_error(format!(
368368
"__str__ returned non-string (type {})",
369369
obj.class().name()

0 commit comments

Comments
 (0)