Reuse unchanged string results in affix removal and partitioning (#9324)

* Reuse unchanged string results in affix removal and partitioning

Assisted-by: Codex:GPT-6

* Qualify the ASCII macro in the string test

Assisted-by: Codex:GPT-6

crates/vm/src/builtins/str.rs

@@ -40,7 +40,6 @@ use itertools::Itertools;
 use memchr::memchr;
 use num_traits::ToPrimitive;
 use rustpython_common::{
-    ascii,
     atomic::{self, PyAtomic, Radium},
     format::{FormatString, FromTemplate},
     hash,
@@ -910,6 +909,16 @@ impl PyStr {
     )
 )]
 impl Py<PyStr> {
+    // `part` is a slice of this string. A full-length slice needs no copy,
+    // except when the operation must return a base str for a subclass.
+    fn result_from_slice(&self, part: &Wtf8, vm: &VirtualMachine) -> PyStrRef {
+        if part.len() == self.byte_len() {
+            PyStr::result_unchanged(self.to_owned(), vm)
+        } else {
+            vm.ctx.new_str(part)
+        }
+    }
+
     #[pymethod]
     #[inline(always)]
     pub fn isascii(&self) -> bool {
@@ -1141,17 +1150,19 @@ impl Py<PyStr> {
     }
 
     #[pymethod]
-    fn removeprefix(&self, prefix: PyStrRef) -> Wtf8Buf {
-        self.as_wtf8()
-            .py_removeprefix(prefix.as_wtf8(), prefix.byte_len(), |s, p| s.starts_with(p))
-            .to_owned()
+    fn removeprefix(&self, prefix: PyStrRef, vm: &VirtualMachine) -> PyStrRef {
+        let part = self
+            .as_wtf8()
+            .py_removeprefix(prefix.as_wtf8(), prefix.byte_len(), |s, p| s.starts_with(p));
+        self.result_from_slice(part, vm)
     }
 
     #[pymethod]
-    fn removesuffix(&self, suffix: PyStrRef) -> Wtf8Buf {
-        self.as_wtf8()
-            .py_removesuffix(suffix.as_wtf8(), suffix.byte_len(), |s, p| s.ends_with(p))
-            .to_owned()
+    fn removesuffix(&self, suffix: PyStrRef, vm: &VirtualMachine) -> PyStrRef {
+        let part = self
+            .as_wtf8()
+            .py_removesuffix(suffix.as_wtf8(), suffix.byte_len(), |s, p| s.ends_with(p));
+        self.result_from_slice(part, vm)
     }
 
     #[pymethod]
@@ -1381,40 +1392,41 @@ impl Py<PyStr> {
 
     #[pymethod]
     pub fn partition(&self, sep: PyStrRef, vm: &VirtualMachine) -> PyResult {
-        let (front, has_mid, back) = self.as_wtf8().py_partition(
-            sep.as_wtf8(),
-            || self.as_wtf8().splitn(2, sep.as_wtf8()),
-            vm,
-        )?;
-        let partition = (
-            self.new_substr(front),
-            if has_mid {
-                sep
-            } else {
-                vm.ctx.new_str(ascii!(""))
-            },
-            self.new_substr(back),
-        );
-        Ok(partition.to_pyobject(vm))
+        self.partition_inner(sep, false, vm)
     }
 
     #[pymethod]
     pub fn rpartition(&self, sep: PyStrRef, vm: &VirtualMachine) -> PyResult {
-        let (back, has_mid, front) = self.as_wtf8().py_partition(
-            sep.as_wtf8(),
-            || self.as_wtf8().rsplitn(2, sep.as_wtf8()),
-            vm,
-        )?;
-        Ok((
-            self.new_substr(front),
-            if has_mid {
-                sep
+        self.partition_inner(sep, true, vm)
+    }
+
+    fn partition_inner(&self, sep: PyStrRef, reverse: bool, vm: &VirtualMachine) -> PyResult {
+        if sep.is_empty() {
+            return Err(vm.new_value_error("empty separator"));
+        }
+        let text = self.as_wtf8();
+        let index = if reverse {
+            text.rfind(sep.as_wtf8())
+        } else {
+            text.find(sep.as_wtf8())
+        };
+        let parts = if let Some(index) = index {
+            (
+                vm.ctx.new_str(&text[..index]),
+                sep.clone(),
+                vm.ctx.new_str(&text[index + sep.byte_len()..]),
+            )
+        } else {
+            let empty = vm.ctx.empty_str.to_owned();
+            // partition preserves the original object, including a subclass,
+            // when there is no separator.
+            if reverse {
+                (empty.clone(), empty, self.to_owned())
             } else {
-                vm.ctx.empty_str.to_owned()
-            },
-            self.new_substr(back),
-        )
-            .to_pyobject(vm))
+                (self.to_owned(), empty.clone(), empty)
+            }
+        };
+        Ok(parts.to_pyobject(vm))
     }
 
     #[pymethod]
@@ -2818,7 +2830,11 @@ mod tests {
                 .unwrap();
             table.set_item("b", vm.ctx.none(), vm).unwrap();
             table
-                .set_item("c", vm.ctx.new_str(ascii!("xda")).into(), vm)
+                .set_item(
+                    "c",
+                    vm.ctx.new_str(rustpython_common::ascii!("xda")).into(),
+                    vm,
+                )
                 .unwrap();
             let translated = Py::<PyStr>::maketrans(
                 table.into(),

extra_tests/snippets/builtin_str.py

@@ -1284,3 +1284,12 @@ def test_percent_callback_order_bounded():
 
 
 test_percent_callback_order_bounded()
+
+
+class PartitionString(str):
+    pass
+
+
+partition_string = PartitionString("abc")
+assert partition_string.partition(":")[0] is partition_string
+assert partition_string.rpartition(":")[2] is partition_string