diff --git a/crates/vm/src/builtins/str.rs b/crates/vm/src/builtins/str.rs index cac1db24e..cf66f093d 100644 --- a/crates/vm/src/builtins/str.rs +++ b/crates/vm/src/builtins/str.rs @@ -893,10 +893,10 @@ impl PyStr { } #[pymethod] - fn strip(&self, args: StripArgs) -> Self { + fn strip(zelf: PyRef, args: StripArgs, vm: &VirtualMachine) -> PyRef { let chars = args.chars; - match self.as_str_kind() { - PyKindStr::Ascii(s) => s + let stripped: &Wtf8 = match zelf.as_str_kind() { + PyKindStr::Ascii(s) if chars.as_ref().is_none_or(|c| c.kind().is_ascii()) => s .py_strip( chars, |s, chars| { @@ -910,21 +910,25 @@ impl PyStr { unsafe { AsciiStr::from_ascii_unchecked(s.as_bytes()) } }, ) + .as_str() .into(), - PyKindStr::Utf8(s) => s + PyKindStr::Utf8(s) if chars.as_ref().is_none_or(|c| c.kind().is_utf8()) => s .py_strip( chars, |s, chars| s.trim_matches(|c| chars.contains(c)), |s| s.trim_matches(unicode::classify::is_space), ) .into(), - PyKindStr::Wtf8(w) => w - .py_strip( - chars, - |s, chars| s.trim_matches(|c| chars.code_points().contains(&c)), - |s| s.trim_matches(|c: CodePoint| c.is_char_and(unicode::classify::is_space)), - ) - .into(), + _ => zelf.as_wtf8().py_strip( + chars, + |s, chars| s.trim_matches(|c| chars.code_points().contains(&c)), + |s| s.trim_matches(|c: CodePoint| c.is_char_and(unicode::classify::is_space)), + ), + }; + if zelf.byte_len() == stripped.len() { + Self::result_unchanged(zelf, vm) + } else { + vm.ctx.new_str(zelf.new_substr(stripped.to_owned())) } } @@ -946,7 +950,7 @@ impl PyStr { |s, chars| s.trim_start_matches(|c| chars.contains_code_point(c)), |s| s.trim_start_matches(|c: CodePoint| c.is_char_and(unicode::classify::is_space)), ); - if s == stripped { + if s.len() == stripped.len() { Self::result_unchanged(zelf, vm) } else { vm.ctx.new_str(stripped) @@ -962,7 +966,7 @@ impl PyStr { |s, chars| s.trim_end_matches(|c| chars.contains_code_point(c)), |s| s.trim_end_matches(|c: CodePoint| c.is_char_and(unicode::classify::is_space)), ); - if s == stripped { + if s.len() == stripped.len() { Self::result_unchanged(zelf, vm) } else { vm.ctx.new_str(stripped) @@ -1110,28 +1114,28 @@ impl PyStr { } #[pymethod] - fn replace(&self, args: ReplaceArgs) -> Wtf8Buf { - use core::cmp::Ordering; - - let s = self.as_wtf8(); + fn replace(zelf: PyRef, args: ReplaceArgs, vm: &VirtualMachine) -> PyRef { let ReplaceArgs { old, new, count } = args; - - match count.cmp(&0) { - Ordering::Less => s.replace(old.as_wtf8(), new.as_wtf8()), - Ordering::Equal => s.to_owned(), - Ordering::Greater => { - let s_is_empty = s.is_empty(); - let old_is_empty = old.is_empty(); - - if s_is_empty && !old_is_empty { - s.to_owned() - } else if s_is_empty && old_is_empty { - new.as_wtf8().to_owned() - } else { - s.replacen(old.as_wtf8(), new.as_wtf8(), count as usize) - } - } + if count == 0 || old.byte_len() > zelf.byte_len() || old.as_wtf8() == new.as_wtf8() { + return Self::result_unchanged(zelf, vm); } + + let s = zelf.as_wtf8(); + let replaced = if count < 0 { + s.replace(old.as_wtf8(), new.as_wtf8()) + } else { + let s_is_empty = s.is_empty(); + let old_is_empty = old.is_empty(); + + if s_is_empty && !old_is_empty { + s.to_owned() + } else if s_is_empty && old_is_empty { + new.as_wtf8().to_owned() + } else { + s.replacen(old.as_wtf8(), new.as_wtf8(), count as usize) + } + }; + vm.ctx.new_str(replaced) } #[pymethod] diff --git a/extra_tests/snippets/builtin_str.py b/extra_tests/snippets/builtin_str.py index 4687d6f5b..c8801a64c 100644 --- a/extra_tests/snippets/builtin_str.py +++ b/extra_tests/snippets/builtin_str.py @@ -110,6 +110,24 @@ s = "^*RustPython*^" assert s.strip("^*") == "RustPython" assert s.lstrip("^*") == "RustPython*^" assert s.rstrip("^*") == "^*RustPython" +assert " abc ".strip(" é") == "abc" +assert "éabcé".strip("é\ud800") == "abc" + + +def test_strip_unchanged(): + class StrSubclass(str): + pass + + for text in ("already clean", "déjà propre", "a\ud800b"): + for method in (str.strip, str.lstrip, str.rstrip): + for chars in (None, "#"): + assert method(text, chars) is text + result = method(StrSubclass(text), chars) + assert result == text + assert type(result) is str + + +test_strip_unchanged() s = "RustPython" assert s.ljust(8) == "RustPython" @@ -1048,6 +1066,27 @@ def test_replace_empty_pattern(): test_replace_empty_pattern() +def test_replace_unchanged(): + for text in ("already clean", "déjà propre", "a\ud800b"): + for old, new, count in ( + (text[:1], "#", 0), + (text[:1], text[:1], -1), + (text + " more", "replacement", -1), + ): + assert text.replace(old, new, count) is text + result = MyString(text).replace(old, new, count) + assert result == text + assert type(result) is str + + # Returning the original string must not bypass argument conversion. + assert_raises(TypeError, lambda: "abc".replace(1, "x", 0)) + assert_raises(TypeError, lambda: "abc".replace("a", 1, 0)) + assert_raises(TypeError, lambda: "abc".replace("a", "a", None)) + + +test_replace_unchanged() + + def test_expandtabs_zero_tabsize(): # With no width to advance to, the tabs come out and nothing else moves. # A tab that followed a character used to ask for a run of usize::MAX