merged from TheAnyKey/p39_string_rem_pre_suffix

TheAnyKey · TheAnyKey · commit 84a6e8e8df72 · 2020-05-08T15:28:46.000Z
diff --git a/tests/snippets/strings.py b/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,73 @@ 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'
+
+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'
+
+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)
+
+
diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
@@ -510,6 +510,22 @@ impl PyString {
             .to_owned()
     }
 
+    #[pymethod]
+    fn removeprefix(&self, pref: PyStringRef) -> PyResult<String> {
+        if self.value.as_str().starts_with(&pref.value) {
+            return Ok(self.value[pref.len()..].to_string());
+        }
+        Ok(self.value.to_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.len()].to_string());
+        }
+        Ok(self.value.to_string())
+    }
+
     #[pymethod]
     fn endswith(&self, args: pystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
         self.value.as_str().py_startsendswith(
