Merge pull request #1897 from youknowone/fix-bytes-find

Fix find/index/count not to crash for bigint start/end arguments

Lib/test/string_tests.py

@@ -156,8 +156,6 @@ class BaseTest:
                     self.assertEqual(rem, 0, '%s != 0 for %s' % (rem, i))
                     self.assertEqual(r1, r2, '%s != %s for %s' % (r1, r2, i))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_find(self):
         self.checkequal(0, 'abcdefghiabc', 'find', 'abc')
         self.checkequal(9, 'abcdefghiabc', 'find', 'abc', 1)
@@ -215,8 +213,6 @@ class BaseTest:
                 if loc != -1:
                     self.assertEqual(i[loc:loc+len(j)], j)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_rfind(self):
         self.checkequal(9,  'abcdefghiabc', 'rfind', 'abc')
         self.checkequal(12, 'abcdefghiabc', 'rfind', '')

Lib/test/test_unicode.py

@@ -199,8 +199,7 @@ class UnicodeTest(string_tests.CommonTest,
         self.checkequal(0, 'a' * 10, 'count', 'a\U00100304')
         self.checkequal(0, '\u0102' * 10, 'count', '\u0102\U00100304')
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
+    @unittest.skip("TODO: RUSTPYTHON")
     def test_find(self):
         string_tests.CommonTest.test_find(self)
         # test implementation details of the memchr fast path
@@ -232,8 +231,7 @@ class UnicodeTest(string_tests.CommonTest,
         self.checkequal(-1, 'a' * 100, 'find', 'a\U00100304')
         self.checkequal(-1, '\u0102' * 100, 'find', '\u0102\U00100304')
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
+    @unittest.skip("TODO: RUSTPYTHON")
     def test_rfind(self):
         string_tests.CommonTest.test_rfind(self)
         # test implementation details of the memrchr fast path

vm/src/obj/objbyteinner.rs

@@ -154,10 +154,10 @@ impl ByteInnerNewOptions {
 pub struct ByteInnerFindOptions {
     #[pyarg(positional_only, optional = false)]
     sub: Either<PyByteInner, PyIntRef>,
-    #[pyarg(positional_only, optional = true)]
-    start: OptionalArg<Option<isize>>,
-    #[pyarg(positional_only, optional = true)]
-    end: OptionalArg<Option<isize>>,
+    #[pyarg(positional_only, default = "None")]
+    start: Option<PyIntRef>,
+    #[pyarg(positional_only, default = "None")]
+    end: Option<PyIntRef>,
 }
 
 impl ByteInnerFindOptions {

vm/src/obj/objstr.rs

@@ -816,63 +816,35 @@ impl PyString {
     }
 
     #[inline]
-    fn _find<F>(
-        &self,
-        sub: PyStringRef,
-        start: OptionalArg<Option<isize>>,
-        end: OptionalArg<Option<isize>>,
-        find: F,
-    ) -> Option<usize>
+    fn _find<F>(&self, args: FindArgs, find: F) -> Option<usize>
     where
         F: Fn(&str, &str) -> Option<usize>,
     {
-        let range = adjust_indices(start, end, self.value.len());
+        let (sub, range) = args.get_value(self.len());
         self.value.py_find(&sub.value, range, find)
     }
 
     #[pymethod]
-    fn find(
-        &self,
-        sub: PyStringRef,
-        start: OptionalArg<Option<isize>>,
-        end: OptionalArg<Option<isize>>,
-    ) -> isize {
-        self._find(sub, start, end, |r, s| r.find(s))
+    fn find(&self, args: FindArgs) -> isize {
+        self._find(args, |r, s| r.find(s))
             .map_or(-1, |v| v as isize)
     }
 
     #[pymethod]
-    fn rfind(
-        &self,
-        sub: PyStringRef,
-        start: OptionalArg<Option<isize>>,
-        end: OptionalArg<Option<isize>>,
-    ) -> isize {
-        self._find(sub, start, end, |r, s| r.rfind(s))
+    fn rfind(&self, args: FindArgs) -> isize {
+        self._find(args, |r, s| r.rfind(s))
             .map_or(-1, |v| v as isize)
     }
 
     #[pymethod]
-    fn index(
-        &self,
-        sub: PyStringRef,
-        start: OptionalArg<Option<isize>>,
-        end: OptionalArg<Option<isize>>,
-        vm: &VirtualMachine,
-    ) -> PyResult<usize> {
-        self._find(sub, start, end, |r, s| r.find(s))
+    fn index(&self, args: FindArgs, vm: &VirtualMachine) -> PyResult<usize> {
+        self._find(args, |r, s| r.find(s))
             .ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
     }
 
     #[pymethod]
-    fn rindex(
-        &self,
-        sub: PyStringRef,
-        start: OptionalArg<Option<isize>>,
-        end: OptionalArg<Option<isize>>,
-        vm: &VirtualMachine,
-    ) -> PyResult<usize> {
-        self._find(sub, start, end, |r, s| r.rfind(s))
+    fn rindex(&self, args: FindArgs, vm: &VirtualMachine) -> PyResult<usize> {
+        self._find(args, |r, s| r.rfind(s))
             .ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
     }
 
@@ -947,15 +919,10 @@ impl PyString {
     }
 
     #[pymethod]
-    fn count(
-        &self,
-        sub: PyStringRef,
-        start: OptionalArg<Option<isize>>,
-        end: OptionalArg<Option<isize>>,
-    ) -> usize {
-        let range = adjust_indices(start, end, self.value.len());
+    fn count(&self, args: FindArgs) -> usize {
+        let (needle, range) = args.get_value(self.len());
         self.value
-            .py_count(&sub.value, range, |h, n| h.matches(n).count())
+            .py_count(&needle.value, range, |h, n| h.matches(n).count())
     }
 
     #[pymethod]
@@ -1256,6 +1223,23 @@ impl TryFromObject for std::ffi::CString {
 
 type SplitArgs = pystr::SplitArgs<PyStringRef, str, char>;
 
+#[derive(FromArgs)]
+pub struct FindArgs {
+    #[pyarg(positional_only, optional = false)]
+    sub: PyStringRef,
+    #[pyarg(positional_only, default = "None")]
+    start: Option<PyIntRef>,
+    #[pyarg(positional_only, default = "None")]
+    end: Option<PyIntRef>,
+}
+
+impl FindArgs {
+    fn get_value(self, len: usize) -> (PyStringRef, std::ops::Range<usize>) {
+        let range = adjust_indices(self.start, self.end, len);
+        (self.sub, range)
+    }
+}
+
 pub fn init(ctx: &PyContext) {
     PyString::extend_class(ctx, &ctx.types.str_type);
 

vm/src/obj/pystr.rs

@@ -1,7 +1,8 @@
 use crate::function::{single_or_tuple_any, OptionalOption};
+use crate::obj::objint::PyIntRef;
 use crate::pyobject::{PyObjectRef, PyResult, TryFromObject, TypeProtocol};
 use crate::vm::VirtualMachine;
-use num_traits::cast::ToPrimitive;
+use num_traits::{cast::ToPrimitive, sign::Signed};
 
 #[derive(FromArgs)]
 pub struct SplitArgs<T, S, E>
@@ -58,10 +59,10 @@ impl ExpandTabsArgs {
 pub struct StartsEndsWithArgs {
     #[pyarg(positional_only, optional = false)]
     affix: PyObjectRef,
-    #[pyarg(positional_only, optional = true)]
-    start: OptionalOption<isize>,
-    #[pyarg(positional_only, optional = true)]
-    end: OptionalOption<isize>,
+    #[pyarg(positional_only, default = "None")]
+    start: Option<PyIntRef>,
+    #[pyarg(positional_only, default = "None")]
+    end: Option<PyIntRef>,
 }
 
 impl StartsEndsWithArgs {
@@ -71,14 +72,25 @@ impl StartsEndsWithArgs {
     }
 }
 
+fn cap_to_isize(py_int: PyIntRef) -> isize {
+    let big = py_int.as_bigint();
+    big.to_isize().unwrap_or_else(|| {
+        if big.is_negative() {
+            std::isize::MIN
+        } else {
+            std::isize::MAX
+        }
+    })
+}
+
 // help get optional string indices
 pub fn adjust_indices(
-    start: OptionalOption<isize>,
-    end: OptionalOption<isize>,
+    start: Option<PyIntRef>,
+    end: Option<PyIntRef>,
     len: usize,
 ) -> std::ops::Range<usize> {
-    let mut start = start.flat_option().unwrap_or(0);
-    let mut end = end.flat_option().unwrap_or(len as isize);
+    let mut start = start.map_or(0, cap_to_isize);
+    let mut end = end.map_or(len as isize, cap_to_isize);
     if end > len as isize {
         end = len as isize;
     } else if end < 0 {