Implement Python 3.9 style string functions: removeprefix and removesuffix PEP616

* implement and test Py39 string operations removeprefix and removesuffix.
* Added test snippets for it using an also contained extension of testutils

Author: TheAnyKey

Lib/test/string_tests.py

@@ -681,6 +681,45 @@ 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')
+        self.checkequal('spam', 'spam', 'removeprefix', 'python')
+        self.checkequal('spam', 'spam', 'removeprefix', 'spider')
+        self.checkequal('spam', 'spam', 'removeprefix', 'spam and eggs')
+
+        self.checkequal('', '', 'removeprefix', '')
+        self.checkequal('', '', 'removeprefix', 'abcde')
+        self.checkequal('abcde', 'abcde', 'removeprefix', '')
+        self.checkequal('', 'abcde', 'removeprefix', 'abcde')
+
+        self.checkraises(TypeError, 'hello', 'removeprefix')
+        self.checkraises(TypeError, 'hello', 'removeprefix', 42)
+        self.checkraises(TypeError, 'hello', 'removeprefix', 42, 'h')
+        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')
+        self.checkequal('spam', 'spam', 'removesuffix', 'python')
+        self.checkequal('spam', 'spam', 'removesuffix', 'blam')
+        self.checkequal('spam', 'spam', 'removesuffix', 'eggs and spam')
+
+        self.checkequal('', '', 'removesuffix', '')
+        self.checkequal('', '', 'removesuffix', 'abcde')
+        self.checkequal('abcde', 'abcde', 'removesuffix', '')
+        self.checkequal('', 'abcde', 'removesuffix', 'abcde')
+
+        self.checkraises(TypeError, 'hello', 'removesuffix')
+        self.checkraises(TypeError, 'hello', 'removesuffix', 42)
+        self.checkraises(TypeError, 'hello', 'removesuffix', 42, 'h')
+        self.checkraises(TypeError, 'hello', 'removesuffix', 'h', 42)
+        self.checkraises(TypeError, 'hello', 'removesuffix', ("lo", "l"))
+
     def test_capitalize(self):
         self.checkequal(' hello ', ' hello ', 'capitalize')
         self.checkequal('Hello ', 'Hello ','capitalize')

tests/snippets/strings.py

@@ -1,4 +1,4 @@
-from testutils import assert_raises, AssertRaises
+from testutils import assert_raises, AssertRaises, skip_if_unsupported
 
 assert "".__eq__(1) == NotImplemented
 assert "a" == 'a'
@@ -471,3 +471,98 @@ def try_mutate_str():
 assert '{:e}'.format(float('inf')) == 'inf'
 assert '{:e}'.format(float('-inf')) == '-inf'
 assert '{:E}'.format(float('inf')) == 'INF'
+
+
+# remove*fix test
+def test_removeprefix():
+    s = 'foobarfoo'
+    s_ref='foobarfoo'
+    assert s.removeprefix('f') == s_ref[1:]
+    assert s.removeprefix('fo') == s_ref[2:]
+    assert s.removeprefix('foo') == s_ref[3:]
+
+    assert s.removeprefix('') == s_ref
+    assert s.removeprefix('bar') == s_ref
+    assert s.removeprefix('lol') == s_ref
+    assert s.removeprefix('_foo') == s_ref
+    assert s.removeprefix('-foo') == s_ref
+    assert s.removeprefix('afoo') == s_ref
+    assert s.removeprefix('*foo') == s_ref
+
+    assert s==s_ref, 'undefined test fail'
+
+    s_uc = '😱foobarfoo🖖'
+    s_ref_uc = '😱foobarfoo🖖'
+    assert s_uc.removeprefix('😱') == s_ref_uc[1:]
+    assert s_uc.removeprefix('😱fo') == s_ref_uc[3:]
+    assert s_uc.removeprefix('😱foo') == s_ref_uc[4:]
+    
+    assert s_uc.removeprefix('🖖') == s_ref_uc
+    assert s_uc.removeprefix('foo') == s_ref_uc
+    assert s_uc.removeprefix(' ') == s_ref_uc
+    assert s_uc.removeprefix('_😱') == s_ref_uc
+    assert s_uc.removeprefix(' 😱') == s_ref_uc
+    assert s_uc.removeprefix('-😱') == s_ref_uc
+    assert s_uc.removeprefix('#😱') == s_ref_uc
+
+def test_removeprefix_types():
+    s='0123456'
+    s_ref='0123456'
+    others=[0,['012']]
+    found=False
+    for o in others:
+        try:
+            s.removeprefix(o)
+        except:
+            found=True
+
+        assert found, f'Removeprefix accepts other type: {type(o)}: {o=}'
+
+def test_removesuffix():
+    s='foobarfoo'
+    s_ref='foobarfoo'
+    assert s.removesuffix('o') == s_ref[:-1]
+    assert s.removesuffix('oo') == s_ref[:-2]
+    assert s.removesuffix('foo') == s_ref[:-3]
+
+    assert s.removesuffix('') == s_ref
+    assert s.removesuffix('bar') == s_ref
+    assert s.removesuffix('lol') == s_ref
+    assert s.removesuffix('foo_') == s_ref
+    assert s.removesuffix('foo-') == s_ref
+    assert s.removesuffix('foo*') == s_ref
+    assert s.removesuffix('fooa') == s_ref
+
+    assert s==s_ref, 'undefined test fail'
+
+    s_uc = '😱foobarfoo🖖'
+    s_ref_uc = '😱foobarfoo🖖'
+    assert s_uc.removesuffix('🖖') == s_ref_uc[:-1]
+    assert s_uc.removesuffix('oo🖖') == s_ref_uc[:-3]
+    assert s_uc.removesuffix('foo🖖') == s_ref_uc[:-4]
+    
+    assert s_uc.removesuffix('😱') == s_ref_uc
+    assert s_uc.removesuffix('foo') == s_ref_uc
+    assert s_uc.removesuffix(' ') == s_ref_uc
+    assert s_uc.removesuffix('🖖_') == s_ref_uc
+    assert s_uc.removesuffix('🖖 ') == s_ref_uc
+    assert s_uc.removesuffix('🖖-') == s_ref_uc
+    assert s_uc.removesuffix('🖖#') == s_ref_uc
+
+def test_removesuffix_types():
+    s='0123456'
+    s_ref='0123456'
+    others=[0,6,['6']]
+    found=False
+    for o in others:
+        try:
+            s.removesuffix(o)
+        except:
+            found=True
+
+        assert found, f'Removesuffix accepts other type: {type(o)}: {o=}'
+
+skip_if_unsupported(3,9,test_removeprefix)
+skip_if_unsupported(3,9,test_removeprefix_types)
+skip_if_unsupported(3,9,test_removesuffix)
+skip_if_unsupported(3,9,test_removesuffix_types)

tests/snippets/testutils.py

@@ -92,4 +92,3 @@ def exec():
         exec()
     else:
         assert False, f'Test cannot performed on this python version. {platform.python_implementation()} {platform.python_version()}'
-

vm/src/obj/objbytearray.rs

@@ -387,6 +387,30 @@ 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 {
+        self.borrow_value().removeprefix(prefix).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 {
+        self.borrow_value().removesuffix(suffix).to_vec().into()
+    }
+
     #[pymethod(name = "split")]
     fn split(&self, options: ByteInnerSplitOptions, vm: &VirtualMachine) -> PyResult {
         self.borrow_value()

vm/src/obj/objbyteinner.rs

@@ -840,6 +840,24 @@ impl PyByteInner {
             .to_vec()
     }
 
+    // new in Python 3.9
+    pub fn removeprefix(&self, prefix: PyByteInner) -> Vec<u8> {
+        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> {
+        self.elements
+            .py_removesuffix(&suffix.elements, suffix.elements.len(), |s, p| {
+                s.ends_with(p)
+            })
+            .to_vec()
+    }
+
     pub fn split<F>(
         &self,
         options: ByteInnerSplitOptions,

vm/src/obj/objbytes.rs

@@ -345,6 +345,30 @@ 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()
+    }
+
     #[pymethod(name = "split")]
     fn split(&self, options: ByteInnerSplitOptions, vm: &VirtualMachine) -> PyResult {
         self.inner

vm/src/obj/objstr.rs

@@ -529,6 +529,36 @@ 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) -> 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) -> String {
+        self.value
+            .as_str()
+            .py_removesuffix(&suff.value, suff.value.len(), |s, p| s.ends_with(p))
+            .to_string()
+    }
+
     #[pymethod]
     fn isalnum(&self) -> bool {
         !self.value.is_empty() && self.value.chars().all(char::is_alphanumeric)

vm/src/obj/pystr.rs

@@ -264,4 +264,37 @@ 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) {
+            self.get_bytes(prefix_len..self.bytes_len())
+        } else {
+            &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) {
+            self.get_bytes(0..self.bytes_len() - suffix_len)
+        } else {
+            &self
+        }
+    }
 }