unified implementation of removeprefix and removesuffix, added pydocs

Author: TheAnyKey

vm/src/obj/objbytearray.rs

@@ -389,24 +389,38 @@ impl PyByteArray {
         self.borrow_value().rstrip(chars).into()
     }
 
+    /// removeprefix($self, prefix, /)
+    ///
+    ///
+    /// Return a bytearray object with the given prefix string removed if present.
+    ///
+    /// If the bytearray starts with the prefix string, return string[len(prefix):]
+    /// Otherwise, return a copy of the original bytearray.
     #[pymethod(name = "removeprefix")]
     fn removeprefix(&self, prefix: PyByteInner) -> PyByteArray {
-        let value = self.borrow_value();
-        if value.elements.starts_with(&prefix.elements) {
-            return value.elements[prefix.elements.len()..].to_vec().into();
-        }
-        value.elements.to_vec().into()
+        self.borrow_value().elements[..]
+            .py_removeprefix(&prefix.elements, prefix.elements.len(), |s, p| {
+                s.starts_with(p)
+            })
+            .to_vec()
+            .into()
     }
 
+    /// removesuffix(self, prefix, /)
+    ///
+    ///
+    /// Return a bytearray object with the given suffix string removed if present.
+    ///
+    /// If the bytearray ends with the suffix string, return string[:len(suffix)]
+    /// Otherwise, return a copy of the original bytearray.
     #[pymethod(name = "removesuffix")]
     fn removesuffix(&self, suffix: PyByteInner) -> PyByteArray {
-        let value = self.borrow_value();
-        if value.elements.ends_with(&suffix.elements) {
-            return value.elements[..value.elements.len() - suffix.elements.len()]
-                .to_vec()
-                .into();
-        }
-        value.elements.to_vec().into()
+        self.borrow_value().elements[..]
+            .py_removesuffix(&suffix.elements, suffix.elements.len(), |s, p| {
+                s.ends_with(p)
+            })
+            .to_vec()
+            .into()
     }
 
     #[pymethod(name = "split")]

vm/src/obj/objbyteinner.rs

@@ -845,18 +845,22 @@ impl PyByteInner {
 
     // new in Python 3.9
     pub fn removeprefix(&self, prefix: PyByteInner) -> Vec<u8> {
-        if self.elements.starts_with(&prefix.elements) {
-            return self.elements[prefix.elements.len()..].to_vec();
-        }
-        self.elements.to_vec()
+        // self.elements.py_removeprefix(&prefix.elements, prefix.elements.len(), |s:&Self, p:&Vec<u8>| s.elements.starts_with(&p)).to_vec()
+
+        self.elements
+            .py_removeprefix(&prefix.elements, prefix.elements.len(), |s, p| {
+                s.starts_with(p)
+            })
+            .to_vec()
     }
 
     // new in Python 3.9
     pub fn removesuffix(&self, suffix: PyByteInner) -> Vec<u8> {
-        if self.elements.ends_with(&suffix.elements) {
-            return self.elements[..self.elements.len() - suffix.elements.len()].to_vec();
-        }
-        self.elements.to_vec()
+        self.elements
+            .py_removesuffix(&suffix.elements, suffix.elements.len(), |s, p| {
+                s.ends_with(p)
+            })
+            .to_vec()
     }
 
     pub fn split<F>(

vm/src/obj/objbytes.rs

@@ -347,11 +347,25 @@ impl PyBytes {
         self.inner.rstrip(chars).into()
     }
 
+    /// removeprefix($self, prefix, /)
+    ///
+    ///
+    /// Return a bytes object with the given prefix string removed if present.
+    ///
+    /// If the bytes starts with the prefix string, return string[len(prefix):]
+    /// Otherwise, return a copy of the original bytes.
     #[pymethod(name = "removeprefix")]
     fn removeprefix(&self, prefix: PyByteInner) -> PyBytes {
         self.inner.removeprefix(prefix).into()
     }
 
+    /// removesuffix(self, prefix, /)
+    ///
+    ///
+    /// Return a bytes object with the given suffix string removed if present.
+    ///
+    /// If the bytes ends with the suffix string, return string[:len(suffix)]
+    /// Otherwise, return a copy of the original bytes.
     #[pymethod(name = "removesuffix")]
     fn removesuffix(&self, suffix: PyByteInner) -> PyBytes {
         self.inner.removesuffix(suffix).into()

vm/src/obj/objstr.rs

@@ -532,20 +532,34 @@ impl PyString {
         )
     }
 
+    /// removeprefix($self, prefix, /)
+    ///
+    ///
+    /// Return a str with the given prefix string removed if present.
+    ///
+    /// If the string starts with the prefix string, return string[len(prefix):]
+    /// Otherwise, return a copy of the original string.
     #[pymethod]
-    fn removeprefix(&self, pref: PyStringRef) -> PyResult<String> {
-        if self.value.as_str().starts_with(&pref.value) {
-            return Ok(self.value[pref.value.len()..].to_string());
-        }
-        Ok(self.value.to_string())
+    fn removeprefix(&self, pref: PyStringRef) -> String {
+        self.value
+            .as_str()
+            .py_removeprefix(&pref.value, pref.value.len(), |s, p| s.starts_with(p))
+            .to_string()
     }
 
+    /// removesuffix(self, prefix, /)
+    ///
+    ///
+    /// Return a str with the given suffix string removed if present.
+    ///
+    /// If the string ends with the suffix string, return string[:len(suffix)]
+    /// Otherwise, return a copy of the original string.
     #[pymethod]
-    fn removesuffix(&self, suff: PyStringRef) -> PyResult<String> {
-        if self.value.as_str().ends_with(&suff.value) {
-            return Ok(self.value[..self.value.len() - suff.value.len()].to_string());
-        }
-        Ok(self.value.to_string())
+    fn removesuffix(&self, suff: PyStringRef) -> String {
+        self.value
+            .as_str()
+            .py_removesuffix(&suff.value, suff.value.len(), |s, p| s.ends_with(p))
+            .to_string()
     }
 
     #[pymethod]

vm/src/obj/pystr.rs

@@ -3,6 +3,7 @@ use crate::obj::objint::PyIntRef;
 use crate::pyobject::{PyObjectRef, PyResult, TryFromObject, TypeProtocol};
 use crate::vm::VirtualMachine;
 use num_traits::{cast::ToPrimitive, sign::Signed};
+use std::ops::Range;
 
 #[derive(FromArgs)]
 pub struct SplitArgs<T, S, E>
@@ -264,4 +265,41 @@ pub trait PyCommonString<E> {
     fn py_rjust(&self, width: usize, fillchar: E) -> Self::Container {
         self.py_pad(width - self.chars_len(), 0, fillchar)
     }
+
+    fn py_removeprefix<FC>(
+        &self,
+        prefix: &Self::Container,
+        prefix_len: usize,
+        is_prefix: FC,
+    ) -> &Self
+    where
+        FC: Fn(&Self, &Self::Container) -> bool,
+    {
+        //if self.py_starts_with(prefix) {
+        if is_prefix(&self, &prefix) {
+            return self.get_bytes(Range {
+                start: prefix_len,
+                end: self.bytes_len(),
+            });
+        }
+        &self
+    }
+
+    fn py_removesuffix<FC>(
+        &self,
+        suffix: &Self::Container,
+        suffix_len: usize,
+        is_suffix: FC,
+    ) -> &Self
+    where
+        FC: Fn(&Self, &Self::Container) -> bool,
+    {
+        if is_suffix(&self, &suffix) {
+            return self.get_bytes(Range {
+                start: 0,
+                end: self.bytes_len() - suffix_len,
+            });
+        }
+        &self
+    }
 }