Make TextIOWrapper wtf8-compatible

coolreader18 · coolreader18 · commit cc6f3d305167 · 2025-03-25T19:05:12.000-05:00
diff --git a/common/src/wtf8/mod.rs b/common/src/wtf8/mod.rs
@@ -613,6 +613,12 @@ impl ToOwned for Wtf8 {
     }
 }
 
+impl PartialEq<str> for Wtf8 {
+    fn eq(&self, other: &str) -> bool {
+        self.as_bytes().eq(other.as_bytes())
+    }
+}
+
 /// Formats the string in double quotes, with characters escaped according to
 /// [`char::escape_debug`] and unpaired surrogates represented as `\u{xxxx}`,
 /// where each `x` is a hexadecimal digit.
@@ -1046,6 +1052,11 @@ impl Wtf8 {
             .strip_suffix(w.as_bytes())
             .map(|w| unsafe { Wtf8::from_bytes_unchecked(w) })
     }
+
+    pub fn replace(&self, from: &Wtf8, to: &Wtf8) -> Wtf8Buf {
+        let w = self.bytes.replace(from, to);
+        unsafe { Wtf8Buf::from_bytes_unchecked(w) }
+    }
 }
 
 impl AsRef<Wtf8> for str {
diff --git a/vm/src/builtins/str.rs b/vm/src/builtins/str.rs
@@ -1408,14 +1408,14 @@ impl PyStrRef {
         (**self).is_empty()
     }
 
-    pub fn concat_in_place(&mut self, other: &str, vm: &VirtualMachine) {
+    pub fn concat_in_place(&mut self, other: &Wtf8, vm: &VirtualMachine) {
         // TODO: call [A]Rc::get_mut on the str to try to mutate the data in place
         if other.is_empty() {
             return;
         }
-        let mut s = String::with_capacity(self.byte_len() + other.len());
-        s.push_str(self.as_ref());
-        s.push_str(other);
+        let mut s = Wtf8Buf::with_capacity(self.byte_len() + other.len());
+        s.push_wtf8(self.as_ref());
+        s.push_wtf8(other);
         *self = PyStr::from(s).into_ref(&vm.ctx);
     }
 }
diff --git a/vm/src/dictdatatype.rs b/vm/src/dictdatatype.rs
@@ -835,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_wtf8() == self.as_ref())
+            Ok(pystr.as_wtf8() == self)
         } else {
             // Fall back to PyObjectRef implementation.
             let s = vm.ctx.new_str(self);
diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs
@@ -129,6 +129,7 @@ mod _io {
             PyMappedThreadMutexGuard, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard,
             PyThreadMutex, PyThreadMutexGuard,
         },
+        common::wtf8::{Wtf8, Wtf8Buf},
         convert::ToPyObject,
         function::{
             ArgBytesLike, ArgIterable, ArgMemoryBuffer, ArgSize, Either, FuncArgs, IntoFuncArgs,
@@ -147,7 +148,6 @@ 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::*},
@@ -1910,10 +1910,12 @@ mod _io {
     impl Newlines {
         /// returns position where the new line starts if found, otherwise position at which to
         /// continue the search after more is read into the buffer
-        fn find_newline(&self, s: &str) -> Result<usize, usize> {
+        fn find_newline(&self, s: &Wtf8) -> Result<usize, usize> {
             let len = s.len();
             match self {
-                Newlines::Universal | Newlines::Lf => s.find('\n').map(|p| p + 1).ok_or(len),
+                Newlines::Universal | Newlines::Lf => {
+                    s.find("\n".as_ref()).map(|p| p + 1).ok_or(len)
+                }
                 Newlines::Passthrough => {
                     let bytes = s.as_bytes();
                     memchr::memchr2(b'\n', b'\r', bytes)
@@ -1928,7 +1930,7 @@ mod _io {
                         })
                         .ok_or(len)
                 }
-                Newlines::Cr => s.find('\n').map(|p| p + 1).ok_or(len),
+                Newlines::Cr => s.find("\n".as_ref()).map(|p| p + 1).ok_or(len),
                 Newlines::Crlf => {
                     // s[searched..] == remaining
                     let mut searched = 0;
@@ -1993,10 +1995,10 @@ mod _io {
             }
         }
 
-        fn len_str(s: &str) -> Self {
+        fn len_str(s: &Wtf8) -> Self {
             Utf8size {
                 bytes: s.len(),
-                chars: s.chars().count(),
+                chars: s.code_points().count(),
             }
         }
     }
@@ -2224,7 +2226,7 @@ mod _io {
 
             let encoding = match args.encoding {
                 None if vm.state.settings.utf8_mode > 0 => PyStr::from("utf-8").into_ref(&vm.ctx),
-                Some(enc) if enc.as_str() != "locale" => enc,
+                Some(enc) if enc.as_wtf8() != "locale" => enc,
                 _ => {
                     // None without utf8_mode or "locale" encoding
                     vm.import("locale", 0)?
@@ -2534,9 +2536,10 @@ mod _io {
                 *snapshot = Some((cookie.dec_flags, input_chunk.clone()));
                 let decoded = vm.call_method(decoder, "decode", (input_chunk, cookie.need_eof))?;
                 let decoded = check_decoded(decoded, vm)?;
-                let pos_is_valid = decoded
-                    .as_str()
-                    .is_char_boundary(cookie.bytes_to_skip as usize);
+                let pos_is_valid = crate::common::wtf8::is_code_point_boundary(
+                    decoded.as_wtf8(),
+                    cookie.bytes_to_skip as usize,
+                );
                 textio.set_decoded_chars(Some(decoded));
                 if !pos_is_valid {
                     return Err(vm.new_os_error("can't restore logical file position".to_owned()));
@@ -2715,9 +2718,9 @@ mod _io {
                 } else if chunks.len() == 1 {
                     chunks.pop().unwrap()
                 } else {
-                    let mut ret = String::with_capacity(chunks_bytes);
+                    let mut ret = Wtf8Buf::with_capacity(chunks_bytes);
                     for chunk in chunks {
-                        ret.push_str(chunk.as_str())
+                        ret.push_wtf8(chunk.as_wtf8())
                     }
                     PyStr::from(ret).into_ref(&vm.ctx)
                 }
@@ -2744,7 +2747,7 @@ mod _io {
 
             let char_len = obj.char_len();
 
-            let data = obj.as_str();
+            let data = obj.as_wtf8();
 
             let replace_nl = match textio.newline {
                 Newlines::Lf => Some("\n"),
@@ -2753,11 +2756,12 @@ mod _io {
                 Newlines::Universal if cfg!(windows) => Some("\r\n"),
                 _ => None,
             };
-            let has_lf = (replace_nl.is_some() || textio.line_buffering) && data.contains('\n');
-            let flush = textio.line_buffering && (has_lf || data.contains('\r'));
+            let has_lf = (replace_nl.is_some() || textio.line_buffering)
+                && data.contains_code_point('\n'.into());
+            let flush = textio.line_buffering && (has_lf || data.contains_code_point('\r'.into()));
             let chunk = if let Some(replace_nl) = replace_nl {
                 if has_lf {
-                    PyStr::from(data.replace('\n', replace_nl)).into_ref(&vm.ctx)
+                    PyStr::from(data.replace("\n".as_ref(), replace_nl.as_ref())).into_ref(&vm.ctx)
                 } else {
                     obj
                 }
@@ -2834,16 +2838,16 @@ mod _io {
                     if self.is_full_slice() {
                         self.0.char_len()
                     } else {
-                        self.slice().chars().count()
+                        self.slice().code_points().count()
                     }
                 }
                 #[inline]
                 fn is_full_slice(&self) -> bool {
                     self.1.len() >= self.0.byte_len()
                 }
                 #[inline]
-                fn slice(&self) -> &str {
-                    &self.0.as_str()[self.1.clone()]
+                fn slice(&self) -> &Wtf8 {
+                    &self.0.as_wtf8()[self.1.clone()]
                 }
                 #[inline]
                 fn slice_pystr(self, vm: &VirtualMachine) -> PyStrRef {
@@ -2894,24 +2898,24 @@ mod _io {
                     Some(remaining) => {
                         assert_eq!(textio.decoded_chars_used.bytes, 0);
                         offset_to_buffer = remaining.utf8_len();
-                        let decoded_chars = decoded_chars.as_str();
+                        let decoded_chars = decoded_chars.as_wtf8();
                         let line = if remaining.is_full_slice() {
                             let mut line = remaining.0;
                             line.concat_in_place(decoded_chars, vm);
                             line
                         } else {
                             let remaining = remaining.slice();
                             let mut s =
-                                String::with_capacity(remaining.len() + decoded_chars.len());
-                            s.push_str(remaining);
-                            s.push_str(decoded_chars);
+                                Wtf8Buf::with_capacity(remaining.len() + decoded_chars.len());
+                            s.push_wtf8(remaining);
+                            s.push_wtf8(decoded_chars);
                             PyStr::from(s).into_ref(&vm.ctx)
                         };
                         start = Utf8size::default();
                         line
                     }
                 };
-                let line_from_start = &line.as_str()[start.bytes..];
+                let line_from_start = &line.as_wtf8()[start.bytes..];
                 let nl_res = textio.newline.find_newline(line_from_start);
                 match nl_res {
                     Ok(p) | Err(p) => {
@@ -2922,7 +2926,7 @@ mod _io {
                                 endpos = start
                                     + Utf8size {
                                         chars: limit - chunked.chars,
-                                        bytes: crate::common::str::char_range_end(
+                                        bytes: crate::common::str::codepoint_range_end(
                                             line_from_start,
                                             limit - chunked.chars,
                                         )
@@ -2963,9 +2967,9 @@ mod _io {
                     chunked += cur_line.byte_len();
                     chunks.push(cur_line);
                 }
-                let mut s = String::with_capacity(chunked);
+                let mut s = Wtf8Buf::with_capacity(chunked);
                 for chunk in chunks {
-                    s.push_str(chunk.slice())
+                    s.push_wtf8(chunk.slice())
                 }
                 PyStr::from(s).into_ref(&vm.ctx)
             } else if let Some(cur_line) = cur_line {
@@ -3100,7 +3104,7 @@ mod _io {
                 return None;
             }
             let decoded_chars = self.decoded_chars.as_ref()?;
-            let avail = &decoded_chars.as_str()[self.decoded_chars_used.bytes..];
+            let avail = &decoded_chars.as_wtf8()[self.decoded_chars_used.bytes..];
             if avail.is_empty() {
                 return None;
             }
@@ -3112,7 +3116,7 @@ mod _io {
                     (PyStr::from(avail).into_ref(&vm.ctx), avail_chars)
                 }
             } else {
-                let s = crate::common::str::get_chars(avail, 0..n);
+                let s = crate::common::str::get_codepoints(avail, 0..n);
                 (PyStr::from(s).into_ref(&vm.ctx), n)
             };
             self.decoded_chars_used += Utf8size {
@@ -3142,11 +3146,11 @@ mod _io {
                 return decoded_chars;
             }
             // TODO: in-place editing of `str` when refcount == 1
-            let decoded_chars_unused = &decoded_chars.as_str()[chars_pos..];
-            let mut s = String::with_capacity(decoded_chars_unused.len() + append_len);
-            s.push_str(decoded_chars_unused);
+            let decoded_chars_unused = &decoded_chars.as_wtf8()[chars_pos..];
+            let mut s = Wtf8Buf::with_capacity(decoded_chars_unused.len() + append_len);
+            s.push_wtf8(decoded_chars_unused);
             if let Some(append) = append {
-                s.push_str(append.as_str())
+                s.push_wtf8(append.as_wtf8())
             }
             PyStr::from(s).into_ref(&vm.ctx)
         }

Original file line number	Diff line number	Diff line change
`@@ -613,6 +613,12 @@ impl ToOwned for Wtf8 {`
`613`	`613`	`}`
`614`	`614`	`}`
`615`	`615`
	`616`	`+impl PartialEq<str> for Wtf8 {`
	`617`	`+ fn eq(&self, other: &str) -> bool {`
	`618`	`+ self.as_bytes().eq(other.as_bytes())`
	`619`	`+ }`
	`620`	`+}`
	`621`	`+`
`616`	`622`	`/// Formats the string in double quotes, with characters escaped according to`
`617`	`623`	/// [`char::escape_debug`] and unpaired surrogates represented as `\u{xxxx}`,
`618`	`624`	/// where each `x` is a hexadecimal digit.
`@@ -1046,6 +1052,11 @@ impl Wtf8 {`
`1046`	`1052`	`.strip_suffix(w.as_bytes())`
`1047`	`1053`	`.map(\|w\| unsafe { Wtf8::from_bytes_unchecked(w) })`
`1048`	`1054`	`}`
	`1055`	`+`
	`1056`	`+ pub fn replace(&self, from: &Wtf8, to: &Wtf8) -> Wtf8Buf {`
	`1057`	`+ let w = self.bytes.replace(from, to);`
	`1058`	`+ unsafe { Wtf8Buf::from_bytes_unchecked(w) }`
	`1059`	`+ }`
`1049`	`1060`	`}`
`1050`	`1061`
`1051`	`1062`	`impl AsRef<Wtf8> for str {`
Original file line number	Diff line number	Diff line change
`@@ -1408,14 +1408,14 @@ impl PyStrRef {`
`1408`	`1408`	`(**self).is_empty()`
`1409`	`1409`	`}`
`1410`	`1410`
`1411`		`- pub fn concat_in_place(&mut self, other: &str, vm: &VirtualMachine) {`
	`1411`	`+ pub fn concat_in_place(&mut self, other: &Wtf8, vm: &VirtualMachine) {`
`1412`	`1412`	`// TODO: call [A]Rc::get_mut on the str to try to mutate the data in place`
`1413`	`1413`	`if other.is_empty() {`
`1414`	`1414`	`return;`
`1415`	`1415`	`}`
`1416`		`- let mut s = String::with_capacity(self.byte_len() + other.len());`
`1417`		`- s.push_str(self.as_ref());`
`1418`		`- s.push_str(other);`
	`1416`	`+ let mut s = Wtf8Buf::with_capacity(self.byte_len() + other.len());`
	`1417`	`+ s.push_wtf8(self.as_ref());`
	`1418`	`+ s.push_wtf8(other);`
`1419`	`1419`	`*self = PyStr::from(s).into_ref(&vm.ctx);`
`1420`	`1420`	`}`
`1421`	`1421`	`}`