Merge pull request RustPython#3164 from moreal/correct-pystr-startswith

Fix `PyStr::startswith`

Author: youknowone

Lib/test/test_htmlparser.py

@@ -266,8 +266,6 @@ def test_get_starttag_text(self):
             ("starttag", "foo:bar", [("one", "1"), ("two", "2")]),
             ("starttag_text", s)])
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_cdata_content(self):
         contents = [
             '<!-- not a comment --> &not-an-entity-ref;',

vm/src/anystr.rs

@@ -61,10 +61,32 @@ pub struct StartsEndsWithArgs {
 }
 
 impl StartsEndsWithArgs {
-    fn get_value(self, len: usize) -> (PyObjectRef, std::ops::Range<usize>) {
-        let range = adjust_indices(self.start, self.end, len);
+    pub fn get_value(self, len: usize) -> (PyObjectRef, Option<std::ops::Range<usize>>) {
+        let range = if self.start.is_some() || self.end.is_some() {
+            Some(adjust_indices(self.start, self.end, len))
+        } else {
+            None
+        };
         (self.affix, range)
     }
+
+    #[inline]
+    pub fn prepare<'s, S, F>(self, s: &'s S, len: usize, substr: F) -> Option<(PyObjectRef, &'s S)>
+    where
+        S: ?Sized + AnyStr<'s>,
+        F: Fn(&'s S, std::ops::Range<usize>) -> &'s S,
+    {
+        let (affix, range) = self.get_value(len);
+        let substr = if let Some(range) = range {
+            if !range.is_normal() {
+                return None;
+            }
+            substr(s, range)
+        } else {
+            s
+        };
+        Some((affix, substr))
+    }
 }
 
 fn saturate_to_isize(py_int: PyIntRef) -> isize {
@@ -145,7 +167,9 @@ pub trait AnyStr<'s>: 's {
     // FIXME: get_chars is expensive for str
     fn get_chars(&self, range: std::ops::Range<usize>) -> &Self;
     fn bytes_len(&self) -> usize;
-    // fn chars_len(&self) -> usize;  // cannot access to cache here
+    // NOTE: str::chars().count() consumes the O(n) time. But pystr::char_len does cache.
+    //       So using chars_len directly is too expensive and the below method shouldn't be implemented.
+    // fn chars_len(&self) -> usize;
     fn is_empty(&self) -> bool;
 
     fn py_add(&self, other: &Self) -> Self::Container {
@@ -191,7 +215,7 @@ pub trait AnyStr<'s>: 's {
     #[inline]
     fn py_startsendswith<T, F>(
         &self,
-        args: StartsEndsWithArgs,
+        affix: PyObjectRef,
         func_name: &str,
         py_type_name: &str,
         func: F,
@@ -201,14 +225,9 @@ pub trait AnyStr<'s>: 's {
         T: TryFromObject,
         F: Fn(&Self, &T) -> bool,
     {
-        let (affix, range) = args.get_value(self.bytes_len());
-        if !range.is_normal() {
-            return Ok(false);
-        }
-        let value = self.get_bytes(range);
         single_or_tuple_any(
             affix,
-            &|s: &T| Ok(func(value, s)),
+            &|s: &T| Ok(func(self, s)),
             &|o| {
                 format!(
                     "{} first arg must be {} or a tuple of {}, not {}",

vm/src/builtins/bytearray.rs

@@ -435,8 +435,14 @@ impl PyByteArray {
 
     #[pymethod]
     fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
-        self.borrow_buf().py_startsendswith(
-            options,
+        let borrowed = self.borrow_buf();
+        let (affix, substr) =
+            match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) {
+                Some(x) => x,
+                None => return Ok(false),
+            };
+        substr.py_startsendswith(
+            affix,
             "endswith",
             "bytes",
             |s, x: &PyBytesInner| s.ends_with(&x.elements[..]),
@@ -450,8 +456,14 @@ impl PyByteArray {
         options: anystr::StartsEndsWithArgs,
         vm: &VirtualMachine,
     ) -> PyResult<bool> {
-        self.borrow_buf().py_startsendswith(
-            options,
+        let borrowed = self.borrow_buf();
+        let (affix, substr) =
+            match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) {
+                Some(x) => x,
+                None => return Ok(false),
+            };
+        substr.py_startsendswith(
+            affix,
             "startswith",
             "bytes",
             |s, x: &PyBytesInner| s.starts_with(&x.elements[..]),

vm/src/builtins/bytes.rs

@@ -275,8 +275,13 @@ impl PyBytes {
 
     #[pymethod]
     fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
-        self.inner.elements[..].py_startsendswith(
-            options,
+        let (affix, substr) =
+            match options.prepare(&self.inner.elements[..], self.len(), |s, r| s.get_bytes(r)) {
+                Some(x) => x,
+                None => return Ok(false),
+            };
+        substr.py_startsendswith(
+            affix,
             "endswith",
             "bytes",
             |s, x: &PyBytesInner| s.ends_with(&x.elements[..]),
@@ -290,8 +295,13 @@ impl PyBytes {
         options: anystr::StartsEndsWithArgs,
         vm: &VirtualMachine,
     ) -> PyResult<bool> {
-        self.inner.elements[..].py_startsendswith(
-            options,
+        let (affix, substr) =
+            match options.prepare(&self.inner.elements[..], self.len(), |s, r| s.get_bytes(r)) {
+                Some(x) => x,
+                None => return Ok(false),
+            };
+        substr.py_startsendswith(
+            affix,
             "startswith",
             "bytes",
             |s, x: &PyBytesInner| s.starts_with(&x.elements[..]),

vm/src/builtins/pystr.rs

@@ -705,9 +705,14 @@ impl PyStr {
     }
 
     #[pymethod]
-    fn endswith(&self, args: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
-        self.as_str().py_startsendswith(
-            args,
+    fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
+        let (affix, substr) =
+            match options.prepare(self.as_str(), self.len(), |s, r| s.get_chars(r)) {
+                Some(x) => x,
+                None => return Ok(false),
+            };
+        substr.py_startsendswith(
+            affix,
             "endswith",
             "str",
             |s, x: &PyStrRef| s.ends_with(x.as_str()),
@@ -716,9 +721,18 @@ impl PyStr {
     }
 
     #[pymethod]
-    fn startswith(&self, args: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
-        self.as_str().py_startsendswith(
-            args,
+    fn startswith(
+        &self,
+        options: anystr::StartsEndsWithArgs,
+        vm: &VirtualMachine,
+    ) -> PyResult<bool> {
+        let (affix, substr) =
+            match options.prepare(self.as_str(), self.len(), |s, r| s.get_chars(r)) {
+                Some(x) => x,
+                None => return Ok(false),
+            };
+        substr.py_startsendswith(
+            affix,
             "startswith",
             "str",
             |s, x: &PyStrRef| s.starts_with(x.as_str()),