removeprefix, suffix: implementation for bytes and bytes array

Author: TheAnyKey

Lib/test/string_tests.py

@@ -681,6 +681,8 @@ def test_replace_overflow(self):
         self.checkraises(OverflowError, A2_16, "replace", "A", A2_16)
         self.checkraises(OverflowError, A2_16, "replace", "AA", A2_16+A2_16)
 
+
+    # Python 3.9
     def test_removeprefix(self):
         self.checkequal('am', 'spam', 'removeprefix', 'sp')
         self.checkequal('spamspam', 'spamspamspam', 'removeprefix', 'spam')
@@ -699,6 +701,7 @@ def test_removeprefix(self):
         self.checkraises(TypeError, 'hello', 'removeprefix', 'h', 42)
         self.checkraises(TypeError, 'hello', 'removeprefix', ("he", "l"))
 
+    # Python 3.9
     def test_removesuffix(self):
         self.checkequal('sp', 'spam', 'removesuffix', 'am')
         self.checkequal('spamspam', 'spamspamspam', 'removesuffix', 'spam')

vm/src/obj/objbytearray.rs

@@ -389,6 +389,26 @@ impl PyByteArray {
         self.borrow_value().rstrip(chars).into()
     }
 
+    #[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()
+    }
+
+    #[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()
+    }
+
     #[pymethod(name = "split")]
     fn split(&self, options: ByteInnerSplitOptions, vm: &VirtualMachine) -> PyResult {
         self.borrow_value()

vm/src/obj/objbyteinner.rs

@@ -843,6 +843,22 @@ impl PyByteInner {
             .to_vec()
     }
 
+    // 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()
+    }
+
+    // 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()
+    }
+
     pub fn split<F>(
         &self,
         options: ByteInnerSplitOptions,

vm/src/obj/objbytes.rs

@@ -347,6 +347,16 @@ impl PyBytes {
         self.inner.rstrip(chars).into()
     }
 
+    #[pymethod(name = "removeprefix")]
+    fn removeprefix(&self, prefix: PyByteInner) -> PyBytes {
+        self.inner.removeprefix(prefix).into()
+    }
+
+    #[pymethod(name = "removesuffix")]
+    fn removesuffix(&self, suffix: PyByteInner) -> PyBytes {
+        self.inner.removesuffix(suffix).into()
+    }
+
     #[pymethod(name = "split")]
     fn split(&self, options: ByteInnerSplitOptions, vm: &VirtualMachine) -> PyResult {
         self.inner