diff --git a/Lib/test/test_cmd_line_script.py b/Lib/test/test_cmd_line_script.py index e40069d78..d5fa1563e 100644 --- a/Lib/test/test_cmd_line_script.py +++ b/Lib/test/test_cmd_line_script.py @@ -574,6 +574,8 @@ class CmdLineTest(unittest.TestCase): self.assertTrue(text[1].startswith(' File ')) self.assertTrue(text[3].startswith('NameError')) + # TODO: RUSTPYTHON + @unittest.expectedFailure def test_non_ascii(self): # Mac OS X denies the creation of a file with an invalid UTF-8 name. # Windows allows creating a name with an arbitrary bytes name, but diff --git a/common/src/encodings.rs b/common/src/encodings.rs index 3aaca68c3..402e0db77 100644 --- a/common/src/encodings.rs +++ b/common/src/encodings.rs @@ -1,5 +1,8 @@ use std::ops::Range; +use num_traits::ToPrimitive; + +use crate::str::StrKind; use crate::wtf8::{Wtf8, Wtf8Buf}; pub type EncodeErrorResult = Result<(EncodeReplace, usize), E>; @@ -7,8 +10,13 @@ pub type EncodeErrorResult = Result<(EncodeReplace, usize), E>; pub type DecodeErrorResult = Result<(S, Option, usize), E>; pub trait StrBuffer: AsRef { - fn is_ascii(&self) -> bool { - self.as_ref().is_ascii() + fn is_compatible_with(&self, kind: StrKind) -> bool { + let s = self.as_ref(); + match kind { + StrKind::Ascii => s.is_ascii(), + StrKind::Utf8 => s.is_utf8(), + StrKind::Wtf8 => true, + } } } @@ -18,7 +26,7 @@ pub trait ErrorHandler { type BytesBuf: AsRef<[u8]>; fn handle_encode_error( &self, - data: &str, + data: &Wtf8, char_range: Range, reason: &str, ) -> EncodeErrorResult; @@ -29,7 +37,7 @@ pub trait ErrorHandler { reason: &str, ) -> DecodeErrorResult; fn error_oob_restart(&self, i: usize) -> Self::Error; - fn error_encoding(&self, data: &str, char_range: Range, reason: &str) -> Self::Error; + fn error_encoding(&self, data: &Wtf8, char_range: Range, reason: &str) -> Self::Error; } pub enum EncodeReplace { Str(S), @@ -118,14 +126,61 @@ where Ok((out, remaining_index)) } +#[inline] +fn encode_utf8_compatible( + s: &Wtf8, + errors: &E, + err_reason: &str, + target_kind: StrKind, +) -> Result, E::Error> { + let full_data = s; + let mut data = s; + let mut char_data_index = 0; + let mut out = Vec::::new(); + while let Some((char_i, (byte_i, _))) = data + .code_point_indices() + .enumerate() + .find(|(_, (_, c))| !target_kind.can_encode(*c)) + { + out.extend_from_slice(&data.as_bytes()[..byte_i]); + let char_start = char_data_index + char_i; + + // number of non-compatible chars between the first non-compatible char and the next compatible char + let non_compat_run_length = data[byte_i..] + .code_points() + .take_while(|c| !target_kind.can_encode(*c)) + .count(); + let char_range = char_start..char_start + non_compat_run_length; + let (replace, char_restart) = + errors.handle_encode_error(full_data, char_range.clone(), err_reason)?; + match replace { + EncodeReplace::Str(s) => { + if s.is_compatible_with(target_kind) { + out.extend_from_slice(s.as_ref().as_bytes()); + } else { + return Err(errors.error_encoding(full_data, char_range, err_reason)); + } + } + EncodeReplace::Bytes(b) => { + out.extend_from_slice(b.as_ref()); + } + } + data = crate::str::try_get_codepoints(full_data, char_restart..) + .ok_or_else(|| errors.error_oob_restart(char_restart))?; + char_data_index = char_restart; + } + out.extend_from_slice(data.as_bytes()); + Ok(out) +} + pub mod utf8 { use super::*; pub const ENCODING_NAME: &str = "utf-8"; #[inline] - pub fn encode(s: &str, _errors: &E) -> Result, E::Error> { - Ok(s.as_bytes().to_vec()) + pub fn encode(s: &Wtf8, errors: &E) -> Result, E::Error> { + encode_utf8_compatible(s, errors, "surrogates not allowed", StrKind::Utf8) } pub fn decode( @@ -175,6 +230,7 @@ pub mod utf8 { } pub mod latin_1 { + use super::*; pub const ENCODING_NAME: &str = "latin-1"; @@ -182,14 +238,14 @@ pub mod latin_1 { const ERR_REASON: &str = "ordinal not in range(256)"; #[inline] - pub fn encode(s: &str, errors: &E) -> Result, E::Error> { + pub fn encode(s: &Wtf8, errors: &E) -> Result, E::Error> { let full_data = s; let mut data = s; let mut char_data_index = 0; let mut out = Vec::::new(); loop { match data - .char_indices() + .code_point_indices() .enumerate() .find(|(_, (_, c))| !c.is_ascii()) { @@ -200,17 +256,16 @@ pub mod latin_1 { Some((char_i, (byte_i, ch))) => { out.extend_from_slice(&data.as_bytes()[..byte_i]); let char_start = char_data_index + char_i; - if (ch as u32) <= 255 { - out.push(ch as u8); - let char_restart = char_start + 1; - data = crate::str::try_get_chars(full_data, char_restart..) - .ok_or_else(|| errors.error_oob_restart(char_restart))?; - char_data_index = char_restart; + if let Some(byte) = ch.to_u32().to_u8() { + out.push(byte); + // if the codepoint is between 128..=255, it's utf8-length is 2 + data = &data[byte_i + 2..]; + char_data_index = char_start + 1; } else { // number of non-latin_1 chars between the first non-latin_1 char and the next latin_1 char let non_latin_1_run_length = data[byte_i..] - .chars() - .take_while(|c| (*c as u32) > 255) + .code_points() + .take_while(|c| c.to_u32() > 255) .count(); let char_range = char_start..char_start + non_latin_1_run_length; let (replace, char_restart) = errors.handle_encode_error( @@ -231,7 +286,7 @@ pub mod latin_1 { out.extend_from_slice(b.as_ref()); } } - data = crate::str::try_get_chars(full_data, char_restart..) + data = crate::str::try_get_codepoints(full_data, char_restart..) .ok_or_else(|| errors.error_oob_restart(char_restart))?; char_data_index = char_restart; } @@ -258,51 +313,8 @@ pub mod ascii { const ERR_REASON: &str = "ordinal not in range(128)"; #[inline] - pub fn encode(s: &str, errors: &E) -> Result, E::Error> { - let full_data = s; - let mut data = s; - let mut char_data_index = 0; - let mut out = Vec::::new(); - loop { - match data - .char_indices() - .enumerate() - .find(|(_, (_, c))| !c.is_ascii()) - { - None => { - out.extend_from_slice(data.as_bytes()); - break; - } - Some((char_i, (byte_i, _))) => { - out.extend_from_slice(&data.as_bytes()[..byte_i]); - let char_start = char_data_index + char_i; - // number of non-ascii chars between the first non-ascii char and the next ascii char - let non_ascii_run_length = - data[byte_i..].chars().take_while(|c| !c.is_ascii()).count(); - let char_range = char_start..char_start + non_ascii_run_length; - let (replace, char_restart) = - errors.handle_encode_error(full_data, char_range.clone(), ERR_REASON)?; - match replace { - EncodeReplace::Str(s) => { - if !s.is_ascii() { - return Err( - errors.error_encoding(full_data, char_range, ERR_REASON) - ); - } - out.extend_from_slice(s.as_ref().as_bytes()); - } - EncodeReplace::Bytes(b) => { - out.extend_from_slice(b.as_ref()); - } - } - data = crate::str::try_get_chars(full_data, char_restart..) - .ok_or_else(|| errors.error_oob_restart(char_restart))?; - char_data_index = char_restart; - continue; - } - } - } - Ok(out) + pub fn encode(s: &Wtf8, errors: &E) -> Result, E::Error> { + encode_utf8_compatible(s, errors, ERR_REASON, StrKind::Ascii) } pub fn decode(data: &[u8], errors: &E) -> Result<(Wtf8Buf, usize), E::Error> { diff --git a/common/src/str.rs b/common/src/str.rs index dbe4e2223..176b5d0f8 100644 --- a/common/src/str.rs +++ b/common/src/str.rs @@ -14,7 +14,7 @@ pub type wchar_t = libc::wchar_t; pub type wchar_t = u32; /// Utf8 + state.ascii (+ PyUnicode_Kind in future) -#[derive(Debug, Copy, Clone, PartialEq)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)] pub enum StrKind { Ascii, Utf8, @@ -41,6 +41,15 @@ impl StrKind { pub fn is_utf8(&self) -> bool { matches!(self, Self::Ascii | Self::Utf8) } + + #[inline(always)] + pub fn can_encode(&self, code: CodePoint) -> bool { + match self { + StrKind::Ascii => code.is_ascii(), + StrKind::Utf8 => code.to_char().is_some(), + StrKind::Wtf8 => true, + } + } } pub trait DeduceStrKind { diff --git a/common/src/wtf8/mod.rs b/common/src/wtf8/mod.rs index 14ff8b5ef..5b31fdc55 100644 --- a/common/src/wtf8/mod.rs +++ b/common/src/wtf8/mod.rs @@ -49,7 +49,7 @@ use std::collections::TryReserveError; use std::string::String; use std::vec::Vec; -use bstr::ByteSlice; +use bstr::{ByteSlice, ByteVec}; mod core_char; mod core_str; @@ -168,6 +168,10 @@ impl CodePoint { pub fn len_wtf8(&self) -> usize { len_utf8(self.value) } + + pub fn is_ascii(&self) -> bool { + self.is_char_and(|c| c.is_ascii()) + } } impl From for CodePoint { @@ -436,6 +440,13 @@ impl Wtf8Buf { self.push_wtf8(code_point.encode_wtf8(&mut [0; MAX_LEN_UTF8])) } + pub fn pop(&mut self) -> Option { + let ch = self.code_points().next_back()?; + let newlen = self.len() - ch.len_wtf8(); + self.bytes.truncate(newlen); + Some(ch) + } + /// Shortens a string to the specified length. /// /// # Panics @@ -448,6 +459,20 @@ impl Wtf8Buf { self.bytes.truncate(new_len) } + /// Inserts a codepoint into this `Wtf8Buf` at a byte position. + #[inline] + pub fn insert(&mut self, idx: usize, c: CodePoint) { + self.insert_wtf8(idx, c.encode_wtf8(&mut [0; MAX_LEN_UTF8])) + } + + /// Inserts a WTF-8 slice into this `Wtf8Buf` at a byte position. + #[inline] + pub fn insert_wtf8(&mut self, idx: usize, w: &Wtf8) { + assert!(is_code_point_boundary(self, idx)); + + self.bytes.insert_str(idx, w) + } + /// Consumes the WTF-8 string and tries to convert it to a vec of bytes. #[inline] pub fn into_bytes(self) -> Vec { @@ -914,6 +939,21 @@ impl Wtf8 { .map(|w| unsafe { Wtf8::from_bytes_unchecked(w) }) } + pub fn trim(&self) -> &Self { + let w = self.bytes.trim(); + unsafe { Wtf8::from_bytes_unchecked(w) } + } + + pub fn trim_start(&self) -> &Self { + let w = self.bytes.trim_start(); + unsafe { Wtf8::from_bytes_unchecked(w) } + } + + pub fn trim_end(&self) -> &Self { + let w = self.bytes.trim_end(); + unsafe { Wtf8::from_bytes_unchecked(w) } + } + pub fn trim_start_matches(&self, f: impl Fn(CodePoint) -> bool) -> &Self { let mut iter = self.code_points(); loop { @@ -958,6 +998,15 @@ impl Wtf8 { memchr::memmem::rfind(self.as_bytes(), pat.as_bytes()) } + pub fn contains(&self, pat: &Wtf8) -> bool { + self.bytes.contains_str(pat) + } + + pub fn contains_code_point(&self, pat: CodePoint) -> bool { + self.bytes + .contains_str(pat.encode_wtf8(&mut [0; MAX_LEN_UTF8])) + } + pub fn get(&self, range: impl ops::RangeBounds) -> Option<&Self> { let start = match range.start_bound() { ops::Bound::Included(&i) => i, @@ -977,6 +1026,26 @@ impl Wtf8 { None } } + + pub fn ends_with(&self, w: &Wtf8) -> bool { + self.bytes.ends_with_str(w) + } + + pub fn starts_with(&self, w: &Wtf8) -> bool { + self.bytes.starts_with_str(w) + } + + pub fn strip_prefix(&self, w: &Wtf8) -> Option<&Self> { + self.bytes + .strip_prefix(w.as_bytes()) + .map(|w| unsafe { Wtf8::from_bytes_unchecked(w) }) + } + + pub fn strip_suffix(&self, w: &Wtf8) -> Option<&Self> { + self.bytes + .strip_suffix(w.as_bytes()) + .map(|w| unsafe { Wtf8::from_bytes_unchecked(w) }) + } } impl AsRef for str { diff --git a/stdlib/src/csv.rs b/stdlib/src/csv.rs index f07b40b3a..4b7913011 100644 --- a/stdlib/src/csv.rs +++ b/stdlib/src/csv.rs @@ -198,9 +198,9 @@ mod _csv { ) -> PyResult { match_class!(match obj.get_attr("lineterminator", vm)? { s @ PyStr => { - Ok(if s.as_str().as_bytes().eq(b"\r\n") { + Ok(if s.as_bytes().eq(b"\r\n") { csv_core::Terminator::CRLF - } else if let Some(t) = s.as_str().as_bytes().first() { + } else if let Some(t) = s.as_bytes().first() { // Due to limitations in the current implementation within csv_core // the support for multiple characters in lineterminator is not complete. // only capture the first character @@ -942,7 +942,7 @@ mod _csv { ), ) })?; - let input = string.as_str().as_bytes(); + let input = string.as_bytes(); if input.is_empty() || input.starts_with(b"\n") { return Ok(PyIterReturn::Return(vm.ctx.new_list(vec![]).into())); } @@ -1101,11 +1101,11 @@ mod _csv { let field: PyObjectRef = field?; let stringified; let data: &[u8] = match_class!(match field { - ref s @ PyStr => s.as_str().as_bytes(), + ref s @ PyStr => s.as_bytes(), crate::builtins::PyNone => b"", ref obj => { stringified = obj.str(vm)?; - stringified.as_str().as_bytes() + stringified.as_bytes() } }); let mut input_offset = 0; diff --git a/stdlib/src/pyexpat.rs b/stdlib/src/pyexpat.rs index 3cfe048f1..2363e6bed 100644 --- a/stdlib/src/pyexpat.rs +++ b/stdlib/src/pyexpat.rs @@ -136,7 +136,7 @@ mod _pyexpat { #[pymethod(name = "Parse")] fn parse(&self, data: PyStrRef, _isfinal: OptionalArg, vm: &VirtualMachine) { - let reader = Cursor::>::new(data.as_str().as_bytes().to_vec()); + let reader = Cursor::>::new(data.as_bytes().to_vec()); let parser = self.create_config().create_reader(reader); self.do_parse(vm, parser); } diff --git a/stdlib/src/pystruct.rs b/stdlib/src/pystruct.rs index 702e5b668..2c7aa1ebf 100644 --- a/stdlib/src/pystruct.rs +++ b/stdlib/src/pystruct.rs @@ -47,7 +47,7 @@ pub(crate) mod _struct { impl IntoStructFormatBytes { fn format_spec(&self, vm: &VirtualMachine) -> PyResult { - FormatSpec::parse(self.0.as_str().as_bytes(), vm) + FormatSpec::parse(self.0.as_bytes(), vm) } } diff --git a/stdlib/src/re.rs b/stdlib/src/re.rs index 5417e03fc..647f4c69a 100644 --- a/stdlib/src/re.rs +++ b/stdlib/src/re.rs @@ -9,10 +9,11 @@ mod re { * system. */ use crate::vm::{ + PyObjectRef, PyPayload, PyResult, VirtualMachine, builtins::{PyInt, PyIntRef, PyStr, PyStrRef}, convert::{ToPyObject, TryFromObject}, function::{OptionalArg, PosArgs}, - match_class, PyObjectRef, PyResult, PyPayload, VirtualMachine, + match_class, }; use num_traits::Signed; use regex::bytes::{Captures, Regex, RegexBuilder}; @@ -158,11 +159,9 @@ mod re { } fn do_sub(pattern: &PyPattern, repl: PyStrRef, search_text: PyStrRef, limit: usize) -> String { - let out = pattern.regex.replacen( - search_text.as_str().as_bytes(), - limit, - repl.as_str().as_bytes(), - ); + let out = pattern + .regex + .replacen(search_text.as_bytes(), limit, repl.as_bytes()); String::from_utf8_lossy(&out).into_owned() } @@ -172,21 +171,21 @@ mod re { regex_text.push_str(pattern.regex.as_str()); let regex = Regex::new(®ex_text).unwrap(); regex - .captures(search_text.as_str().as_bytes()) + .captures(search_text.as_bytes()) .map(|captures| create_match(search_text.clone(), captures)) } fn do_search(regex: &PyPattern, search_text: PyStrRef) -> Option { regex .regex - .captures(search_text.as_str().as_bytes()) + .captures(search_text.as_bytes()) .map(|captures| create_match(search_text.clone(), captures)) } fn do_findall(vm: &VirtualMachine, pattern: &PyPattern, search_text: PyStrRef) -> PyResult { let out = pattern .regex - .captures_iter(search_text.as_str().as_bytes()) + .captures_iter(search_text.as_bytes()) .map(|captures| match captures.len() { 1 => { let full = captures.get(0).unwrap().as_bytes(); @@ -232,7 +231,7 @@ mod re { .map(|i| i.try_to_primitive::(vm)) .transpose()? .unwrap_or(0); - let text = search_text.as_str().as_bytes(); + let text = search_text.as_bytes(); // essentially Regex::split, but it outputs captures as well let mut output = Vec::new(); let mut last = 0; @@ -332,9 +331,7 @@ mod re { #[pymethod] fn sub(&self, repl: PyStrRef, text: PyStrRef, vm: &VirtualMachine) -> PyResult { - let replaced_text = self - .regex - .replace_all(text.as_str().as_bytes(), repl.as_str().as_bytes()); + let replaced_text = self.regex.replace_all(text.as_bytes(), repl.as_bytes()); let replaced_text = String::from_utf8_lossy(&replaced_text).into_owned(); Ok(vm.ctx.new_str(replaced_text)) } diff --git a/stdlib/src/sqlite.rs b/stdlib/src/sqlite.rs index 85ce8d80f..d20324b5e 100644 --- a/stdlib/src/sqlite.rs +++ b/stdlib/src/sqlite.rs @@ -3000,7 +3000,7 @@ mod _sqlite { ) -> PyResult<*const libc::c_char> { BEGIN_STATEMENTS .iter() - .find(|&&x| x[6..].eq_ignore_ascii_case(s.as_str().as_bytes())) + .find(|&&x| x[6..].eq_ignore_ascii_case(s.as_bytes())) .map(|&x| x.as_ptr().cast()) .ok_or_else(|| { vm.new_value_error( diff --git a/stdlib/src/ssl.rs b/stdlib/src/ssl.rs index 2b8ffc7d8..9bfdf883c 100644 --- a/stdlib/src/ssl.rs +++ b/stdlib/src/ssl.rs @@ -708,7 +708,7 @@ mod _ssl { if !s.is_ascii() { return Err(invalid_cadata(vm)); } - X509::stack_from_pem(s.as_str().as_bytes()) + X509::stack_from_pem(s.as_bytes()) } Either::B(b) => b.with_ref(x509_stack_from_der), }; diff --git a/vm/src/builtins/str.rs b/vm/src/builtins/str.rs index 75921bac7..e8c19c39e 100644 --- a/vm/src/builtins/str.rs +++ b/vm/src/builtins/str.rs @@ -429,6 +429,10 @@ impl PyStr { self.data.as_str_kind() } + pub fn is_utf8(&self) -> bool { + self.kind().is_utf8() + } + fn char_all(&self, test: F) -> bool where F: Fn(char) -> bool, @@ -725,7 +729,7 @@ impl PyStr { .py_strip( chars, |s, chars| s.trim_matches(|c| chars.code_points().contains(&c)), - |s| s.trim_matches(|c| c.is_char_and(char::is_whitespace)), + |s| s.trim(), ) .into(), } @@ -737,10 +741,10 @@ impl PyStr { chars: OptionalOption, vm: &VirtualMachine, ) -> PyRef { - let s = zelf.as_str(); + let s = zelf.as_wtf8(); let stripped = s.py_strip( chars, - |s, chars| s.trim_start_matches(|c| chars.contains(c)), + |s, chars| s.trim_start_matches(|c| chars.contains_code_point(c)), |s| s.trim_start(), ); if s == stripped { @@ -756,10 +760,10 @@ impl PyStr { chars: OptionalOption, vm: &VirtualMachine, ) -> PyRef { - let s = zelf.as_str(); + let s = zelf.as_wtf8(); let stripped = s.py_strip( chars, - |s, chars| s.trim_end_matches(|c| chars.contains(c)), + |s, chars| s.trim_end_matches(|c| chars.contains_code_point(c)), |s| s.trim_end(), ); if s == stripped { @@ -772,7 +776,7 @@ impl PyStr { #[pymethod] fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult { let (affix, substr) = - match options.prepare(self.as_str(), self.len(), |s, r| s.get_chars(r)) { + match options.prepare(self.as_wtf8(), self.len(), |s, r| s.get_chars(r)) { Some(x) => x, None => return Ok(false), }; @@ -780,7 +784,7 @@ impl PyStr { &affix, "endswith", "str", - |s, x: &Py| s.ends_with(x.as_str()), + |s, x: &Py| s.ends_with(x.as_wtf8()), vm, ) } @@ -792,7 +796,7 @@ impl PyStr { vm: &VirtualMachine, ) -> PyResult { let (affix, substr) = - match options.prepare(self.as_str(), self.len(), |s, r| s.get_chars(r)) { + match options.prepare(self.as_wtf8(), self.len(), |s, r| s.get_chars(r)) { Some(x) => x, None => return Ok(false), }; @@ -800,7 +804,7 @@ impl PyStr { &affix, "startswith", "str", - |s, x: &Py| s.starts_with(x.as_str()), + |s, x: &Py| s.starts_with(x.as_wtf8()), vm, ) } diff --git a/vm/src/builtins/type.rs b/vm/src/builtins/type.rs index e8ad67a66..969d6db93 100644 --- a/vm/src/builtins/type.rs +++ b/vm/src/builtins/type.rs @@ -768,7 +768,7 @@ impl PyType { value.class().slot_name(), )) })?; - if name.as_str().as_bytes().contains(&0) { + if name.as_bytes().contains(&0) { return Err(vm.new_value_error("type name must not contain null characters".to_owned())); } @@ -811,7 +811,7 @@ impl Constructor for PyType { let (name, bases, dict, kwargs): (PyStrRef, PyTupleRef, PyDictRef, KwArgs) = args.clone().bind(vm)?; - if name.as_str().as_bytes().contains(&0) { + if name.as_bytes().contains(&0) { return Err(vm.new_value_error("type name must not contain null characters".to_owned())); } diff --git a/vm/src/function/buffer.rs b/vm/src/function/buffer.rs index 91379e7a7..80b36833e 100644 --- a/vm/src/function/buffer.rs +++ b/vm/src/function/buffer.rs @@ -152,7 +152,7 @@ impl ArgStrOrBytesLike { pub fn borrow_bytes(&self) -> BorrowedValue<'_, [u8]> { match self { Self::Buf(b) => b.borrow_buf(), - Self::Str(s) => s.as_str().as_bytes().into(), + Self::Str(s) => s.as_bytes().into(), } } } @@ -195,7 +195,7 @@ impl ArgAsciiBuffer { #[inline] pub fn with_ref(&self, f: impl FnOnce(&[u8]) -> R) -> R { match self { - Self::String(s) => f(s.as_str().as_bytes()), + Self::String(s) => f(s.as_bytes()), Self::Buffer(buffer) => buffer.with_ref(f), } } diff --git a/vm/src/function/fspath.rs b/vm/src/function/fspath.rs index 69f11eb65..ab5cd093b 100644 --- a/vm/src/function/fspath.rs +++ b/vm/src/function/fspath.rs @@ -26,7 +26,7 @@ impl FsPath { let match1 = |obj: PyObjectRef| { let pathlike = match_class!(match obj { s @ PyStr => { - check_nul(s.as_str().as_bytes())?; + check_nul(s.as_bytes())?; FsPath::Str(s) } b @ PyBytes => { @@ -61,7 +61,11 @@ impl FsPath { pub fn as_os_str(&self, vm: &VirtualMachine) -> PyResult<&OsStr> { // TODO: FS encodings match self { - FsPath::Str(s) => Ok(s.as_str().as_ref()), + FsPath::Str(s) => { + // XXX RUSTPYTHON: this is sketchy on windows; it's not guaranteed that its + // OsStr encoding will always be compatible with WTF-8. + Ok(unsafe { OsStr::from_encoded_bytes_unchecked(s.as_wtf8().as_bytes()) }) + } FsPath::Bytes(b) => Self::bytes_as_osstr(b.as_bytes(), vm), } } @@ -69,7 +73,7 @@ impl FsPath { pub fn as_bytes(&self) -> &[u8] { // TODO: FS encodings match self { - FsPath::Str(s) => s.as_str().as_bytes(), + FsPath::Str(s) => s.as_bytes(), FsPath::Bytes(b) => b.as_bytes(), } } diff --git a/vm/src/stdlib/codecs.rs b/vm/src/stdlib/codecs.rs index 9f58c8693..9018ad0e3 100644 --- a/vm/src/stdlib/codecs.rs +++ b/vm/src/stdlib/codecs.rs @@ -2,9 +2,9 @@ pub(crate) use _codecs::make_module; #[pymodule] mod _codecs { - use rustpython_common::wtf8::Wtf8Buf; - use crate::common::encodings; + use crate::common::str::StrKind; + use crate::common::wtf8::{Wtf8, Wtf8Buf}; use crate::{ AsObject, PyObject, PyObjectRef, PyResult, TryFromBorrowedObject, VirtualMachine, builtins::{PyBaseExceptionRef, PyBytes, PyBytesRef, PyStr, PyStrRef, PyTuple}, @@ -105,8 +105,8 @@ mod _codecs { } } impl encodings::StrBuffer for PyStrRef { - fn is_ascii(&self) -> bool { - PyStr::is_ascii(self) + fn is_compatible_with(&self, kind: StrKind) -> bool { + self.kind() <= kind } } impl encodings::ErrorHandler for ErrorsHandler<'_> { @@ -116,7 +116,7 @@ mod _codecs { fn handle_encode_error( &self, - data: &str, + data: &Wtf8, char_range: Range, reason: &str, ) -> PyResult<(encodings::EncodeReplace, usize)> { @@ -219,7 +219,7 @@ mod _codecs { fn error_encoding( &self, - data: &str, + data: &Wtf8, char_range: Range, reason: &str, ) -> Self::Error { @@ -251,10 +251,10 @@ mod _codecs { #[inline] fn encode<'a, F>(self, name: &'a str, encode: F, vm: &'a VirtualMachine) -> EncodeResult where - F: FnOnce(&str, &ErrorsHandler<'a>) -> PyResult>, + F: FnOnce(&Wtf8, &ErrorsHandler<'a>) -> PyResult>, { let errors = ErrorsHandler::new(name, self.errors, vm); - let encoded = encode(self.s.as_str(), &errors)?; + let encoded = encode(self.s.as_wtf8(), &errors)?; Ok((encoded, self.s.char_len())) } } @@ -312,6 +312,9 @@ mod _codecs { #[pyfunction] fn utf_8_encode(args: EncodeArgs, vm: &VirtualMachine) -> EncodeResult { + if args.s.is_utf8() { + return Ok((args.s.as_bytes().to_vec(), args.s.byte_len())); + } do_codec!(utf8::encode, args, vm) } @@ -323,7 +326,7 @@ mod _codecs { #[pyfunction] fn latin_1_encode(args: EncodeArgs, vm: &VirtualMachine) -> EncodeResult { if args.s.is_ascii() { - return Ok((args.s.as_str().as_bytes().to_vec(), args.s.byte_len())); + return Ok((args.s.as_bytes().to_vec(), args.s.byte_len())); } do_codec!(latin_1::encode, args, vm) } @@ -336,7 +339,7 @@ mod _codecs { #[pyfunction] fn ascii_encode(args: EncodeArgs, vm: &VirtualMachine) -> EncodeResult { if args.s.is_ascii() { - return Ok((args.s.as_str().as_bytes().to_vec(), args.s.byte_len())); + return Ok((args.s.as_bytes().to_vec(), args.s.byte_len())); } do_codec!(ascii::encode, args, vm) } diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs index 0f472fd94..84dddee87 100644 --- a/vm/src/stdlib/io.rs +++ b/vm/src/stdlib/io.rs @@ -147,6 +147,7 @@ mod _io { use crossbeam_utils::atomic::AtomicCell; use malachite_bigint::{BigInt, BigUint}; use num_traits::ToPrimitive; + use rustpython_common::wtf8::Wtf8Buf; use std::{ borrow::Cow, io::{self, Cursor, SeekFrom, prelude::*}, @@ -2082,7 +2083,7 @@ mod _io { impl PendingWrite { fn as_bytes(&self) -> &[u8] { match self { - Self::Utf8(s) => s.as_str().as_bytes(), + Self::Utf8(s) => s.as_bytes(), Self::Bytes(b) => b.as_bytes(), } } @@ -3305,14 +3306,14 @@ mod _io { }; let orig_output: PyStrRef = output.try_into_value(vm)?; // this being Cow::Owned means we need to allocate a new string - let mut output = Cow::Borrowed(orig_output.as_str()); + let mut output = Cow::Borrowed(orig_output.as_wtf8()); if self.pendingcr && (final_ || !output.is_empty()) { - output = ["\r", &*output].concat().into(); + output.to_mut().insert(0, '\r'.into()); self.pendingcr = false; } if !final_ { - if let Some(s) = output.strip_suffix('\r') { - output = s.to_owned().into(); + if let Some(s) = output.strip_suffix("\r".as_ref()) { + output = Cow::Owned(s.to_owned()); self.pendingcr = true; } } @@ -3321,19 +3322,21 @@ mod _io { return Ok(vm.ctx.empty_str.to_owned()); } - if (self.seennl == SeenNewline::LF || self.seennl.is_empty()) && !output.contains('\r') + if (self.seennl == SeenNewline::LF || self.seennl.is_empty()) + && !output.contains_code_point('\r'.into()) { - if self.seennl.is_empty() && output.contains('\n') { + if self.seennl.is_empty() && output.contains_code_point('\n'.into()) { self.seennl.insert(SeenNewline::LF); } } else if !self.translate { - let mut matches = output.match_indices(['\r', '\n']); + let output = output.as_bytes(); + let mut matches = memchr::memchr2_iter(b'\r', b'\n', output); while !self.seennl.is_all() { - let Some((i, c)) = matches.next() else { break }; - match c { - "\n" => self.seennl.insert(SeenNewline::LF), + let Some(i) = matches.next() else { break }; + match output[i] { + b'\n' => self.seennl.insert(SeenNewline::LF), // if c isn't \n, it can only be \r - _ if output[i + 1..].starts_with('\n') => { + _ if output.get(i + 1) == Some(&b'\n') => { matches.next(); self.seennl.insert(SeenNewline::CRLF); } @@ -3341,30 +3344,31 @@ mod _io { } } } else { - let mut chunks = output.match_indices(['\r', '\n']); - let mut new_string = String::with_capacity(output.len()); + let bytes = output.as_bytes(); + let mut matches = memchr::memchr2_iter(b'\r', b'\n', bytes); + let mut new_string = Wtf8Buf::with_capacity(output.len()); let mut last_modification_index = 0; - while let Some((cr_index, chunk)) = chunks.next() { - if chunk == "\r" { + while let Some(cr_index) = matches.next() { + if bytes[cr_index] == b'\r' { // skip copying the CR let mut next_chunk_index = cr_index + 1; - if output[cr_index + 1..].starts_with('\n') { - chunks.next(); + if bytes.get(cr_index + 1) == Some(&b'\n') { + matches.next(); self.seennl.insert(SeenNewline::CRLF); // skip the LF too next_chunk_index += 1; } else { self.seennl.insert(SeenNewline::CR); } - new_string.push_str(&output[last_modification_index..cr_index]); - new_string.push('\n'); + new_string.push_wtf8(&output[last_modification_index..cr_index]); + new_string.push_char('\n'); last_modification_index = next_chunk_index; } else { self.seennl.insert(SeenNewline::LF); } } - new_string.push_str(&output[last_modification_index..]); - output = new_string.into(); + new_string.push_wtf8(&output[last_modification_index..]); + output = Cow::Owned(new_string); } Ok(match output { @@ -3404,7 +3408,7 @@ mod _io { ) -> PyResult { let raw_bytes = object .flatten() - .map_or_else(Vec::new, |v| v.as_str().as_bytes().to_vec()); + .map_or_else(Vec::new, |v| v.as_bytes().to_vec()); StringIO { buffer: PyRwLock::new(BufferedIO::new(Cursor::new(raw_bytes))), @@ -3453,7 +3457,7 @@ mod _io { // write string to underlying vector #[pymethod] fn write(&self, data: PyStrRef, vm: &VirtualMachine) -> PyResult { - let bytes = data.as_str().as_bytes(); + let bytes = data.as_bytes(); self.buffer(vm)? .write(bytes) .ok_or_else(|| vm.new_type_error("Error Writing String".to_owned())) diff --git a/vm/src/stdlib/operator.rs b/vm/src/stdlib/operator.rs index d1a4b376e..d8ff1715f 100644 --- a/vm/src/stdlib/operator.rs +++ b/vm/src/stdlib/operator.rs @@ -328,7 +328,7 @@ mod _operator { "comparing strings with non-ASCII characters is not supported".to_owned(), )); } - cmp::timing_safe_cmp(a.as_str().as_bytes(), b.as_str().as_bytes()) + cmp::timing_safe_cmp(a.as_bytes(), b.as_bytes()) } (Either::B(a), Either::B(b)) => { a.with_ref(|a| b.with_ref(|b| cmp::timing_safe_cmp(a, b))) diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs index 39701cb3a..9b196b6d0 100644 --- a/vm/src/stdlib/os.rs +++ b/vm/src/stdlib/os.rs @@ -376,7 +376,7 @@ pub(super) mod _os { fn env_bytes_as_bytes(obj: &Either) -> &[u8] { match obj { - Either::A(s) => s.as_str().as_bytes(), + Either::A(s) => s.as_bytes(), Either::B(b) => b.as_bytes(), } }