Refactor and fix str/bytes count/find/index

This commit is contained in:
Jeong YunWon
2020-04-08 01:15:10 +09:00
parent f5de59a4bd
commit e28ab72b4e
5 changed files with 103 additions and 136 deletions

View File

@@ -606,8 +606,6 @@ class BaseBytesTest:
ValueError, r'byte must be in range\(0, 256\)',
b.find, index)
# TODO: RUSTPYTHON
@unittest.expectedFailure
def test_rfind(self):
b = self.type2test(b'mississippi')
i = 105
@@ -647,8 +645,6 @@ class BaseBytesTest:
self.assertEqual(b.index(i, 1, 3), 1)
self.assertRaises(ValueError, b.index, w, 1, 3)
# TODO: RUSTPYTHON
@unittest.expectedFailure
def test_rindex(self):
b = self.type2test(b'mississippi')
i = 105

View File

@@ -324,30 +324,26 @@ impl PyByteArray {
#[pymethod(name = "find")]
fn find(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
self.inner.borrow().find(options, false, vm)
let index = self.inner.borrow().find(options, false, vm)?;
Ok(index.map_or(-1, |v| v as isize))
}
#[pymethod(name = "index")]
fn index(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
let res = self.inner.borrow().find(options, false, vm)?;
if res == -1 {
return Err(vm.new_value_error("substring not found".to_owned()));
}
Ok(res)
fn index(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
let index = self.inner.borrow().find(options, false, vm)?;
index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
}
#[pymethod(name = "rfind")]
fn rfind(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
self.inner.borrow().find(options, true, vm)
let index = self.inner.borrow().find(options, true, vm)?;
Ok(index.map_or(-1, |v| v as isize))
}
#[pymethod(name = "rindex")]
fn rindex(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
let res = self.inner.borrow().find(options, true, vm)?;
if res == -1 {
return Err(vm.new_value_error("substring not found".to_owned()));
}
Ok(res)
fn rindex(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
let index = self.inner.borrow().find(options, true, vm)?;
index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
}
#[pymethod(name = "remove")]

View File

@@ -13,7 +13,7 @@ use super::objmemory::PyMemoryView;
use super::objnone::PyNoneRef;
use super::objsequence::{is_valid_slice_arg, PySliceableSequence};
use super::objslice::PySliceRef;
use super::objstr::{self, PyString, PyStringRef};
use super::objstr::{self, adjust_indices, PyString, PyStringRef, StringRange};
use super::objtuple::PyTupleRef;
use crate::function::OptionalArg;
use crate::pyhash;
@@ -137,34 +137,22 @@ pub struct ByteInnerFindOptions {
#[pyarg(positional_only, optional = false)]
sub: Either<PyByteInner, PyIntRef>,
#[pyarg(positional_only, optional = true)]
start: OptionalArg<Option<PyIntRef>>,
start: OptionalArg<Option<isize>>,
#[pyarg(positional_only, optional = true)]
end: OptionalArg<Option<PyIntRef>>,
end: OptionalArg<Option<isize>>,
}
impl ByteInnerFindOptions {
pub fn get_value(
self,
elements: &[u8],
len: usize,
vm: &VirtualMachine,
) -> PyResult<(Vec<u8>, Range<usize>)> {
) -> PyResult<(Vec<u8>, std::ops::Range<usize>)> {
let sub = match self.sub {
Either::A(v) => v.elements.to_vec(),
Either::B(int) => vec![int.as_bigint().byte_or(vm)?],
};
let start = match self.start {
OptionalArg::Present(Some(int)) => Some(int.as_bigint().clone()),
_ => None,
};
let end = match self.end {
OptionalArg::Present(Some(int)) => Some(int.as_bigint().clone()),
_ => None,
};
let range = elements.to_vec().get_slice_range(&start, &end);
let range = adjust_indices(self.start, self.end, len);
Ok((sub, range))
}
}
@@ -808,25 +796,18 @@ impl PyByteInner {
}
pub fn count(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
let (sub, range) = options.get_value(&self.elements, vm)?;
if sub.is_empty() {
return Ok(self.len() + 1);
let (needle, range) = options.get_value(self.elements.len(), vm)?;
if !range.is_normal() {
return Ok(0);
}
let mut total: usize = 0;
let mut i_start = range.start;
let i_end = range.end;
for i in self.elements.do_slice(range) {
if i_start + sub.len() <= i_end
&& i == sub[0]
&& &self.elements[i_start..(i_start + sub.len())] == sub.as_slice()
{
total += 1;
}
i_start += 1;
if needle.is_empty() {
return Ok(range.len() + 1);
}
let haystack = &self.elements[range];
let total = haystack
.windows(needle.len())
.filter(|w| *w == needle.as_slice())
.count();
Ok(total)
}
@@ -884,37 +865,36 @@ impl PyByteInner {
Ok(suff.as_slice() == &self.elements.do_slice(range)[offset])
}
#[inline]
pub fn find(
&self,
options: ByteInnerFindOptions,
reverse: bool,
vm: &VirtualMachine,
) -> PyResult<isize> {
let (sub, range) = options.get_value(&self.elements, vm)?;
// not allowed for this method
if range.end < range.start {
return Ok(-1isize);
) -> PyResult<Option<usize>> {
let (needle, range) = options.get_value(self.elements.len(), vm)?;
if !range.is_normal() {
return Ok(None);
}
let start = range.start;
let end = range.end;
if needle.is_empty() {
return Ok(Some(if reverse { range.end } else { range.start }));
}
let haystack = &self.elements[range.clone()];
let windows = haystack.windows(needle.len());
if reverse {
let slice = self.elements.do_slice_reverse(range);
for (n, _) in slice.iter().enumerate() {
if n + sub.len() <= slice.len() && &slice[n..n + sub.len()] == sub.as_slice() {
return Ok((end - n - 1) as isize);
for (i, w) in windows.rev().enumerate() {
if w == needle.as_slice() {
return Ok(Some(range.end - i - needle.len()));
}
}
} else {
let slice = self.elements.do_slice(range);
for (n, _) in slice.iter().enumerate() {
if n + sub.len() <= slice.len() && &slice[n..n + sub.len()] == sub.as_slice() {
return Ok((start + n) as isize);
for (i, w) in windows.enumerate() {
if w == needle.as_slice() {
return Ok(Some(range.start + i));
}
}
};
Ok(-1isize)
}
Ok(None)
}
pub fn maketrans(from: PyByteInner, to: PyByteInner, vm: &VirtualMachine) -> PyResult {

View File

@@ -294,30 +294,26 @@ impl PyBytes {
#[pymethod(name = "find")]
fn find(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
self.inner.find(options, false, vm)
let index = self.inner.find(options, false, vm)?;
Ok(index.map_or(-1, |v| v as isize))
}
#[pymethod(name = "index")]
fn index(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
let res = self.inner.find(options, false, vm)?;
if res == -1 {
return Err(vm.new_value_error("substring not found".to_owned()));
}
Ok(res)
fn index(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
let index = self.inner.find(options, false, vm)?;
index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
}
#[pymethod(name = "rfind")]
fn rfind(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
self.inner.find(options, true, vm)
let index = self.inner.find(options, true, vm)?;
Ok(index.map_or(-1, |v| v as isize))
}
#[pymethod(name = "rindex")]
fn rindex(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
let res = self.inner.find(options, true, vm)?;
if res == -1 {
return Err(vm.new_value_error("substring not found".to_owned()));
}
Ok(res)
fn rindex(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
let index = self.inner.find(options, true, vm)?;
index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
}
#[pymethod(name = "translate")]

View File

@@ -579,8 +579,9 @@ impl PyString {
end: OptionalArg<Option<isize>>,
vm: &VirtualMachine,
) -> PyResult<bool> {
if let Some((start, end)) = adjust_indices(start, end, self.value.len()) {
let value = &self.value[start..end];
let range = adjust_indices(start, end, self.value.len());
if range.is_normal() {
let value = &self.value[range];
single_or_tuple_any(
suffix,
|s: PyStringRef| Ok(value.ends_with(&s.value)),
@@ -605,8 +606,9 @@ impl PyString {
end: OptionalArg<Option<isize>>,
vm: &VirtualMachine,
) -> PyResult<bool> {
if let Some((start, end)) = adjust_indices(start, end, self.value.len()) {
let value = &self.value[start..end];
let range = adjust_indices(start, end, self.value.len());
if range.is_normal() {
let value = &self.value[range];
single_or_tuple_any(
prefix,
|s: PyStringRef| Ok(value.starts_with(&s.value)),
@@ -898,6 +900,25 @@ impl PyString {
Ok(joined)
}
fn _find<F>(
&self,
sub: PyStringRef,
start: OptionalArg<Option<isize>>,
end: OptionalArg<Option<isize>>,
find: F,
) -> Option<usize>
where
F: Fn(&str, &str) -> Option<usize>,
{
let range = adjust_indices(start, end, self.value.len());
if range.is_normal() {
if let Some(index) = find(&self.value[range.clone()], &sub.value) {
return Some(range.start + index);
}
}
None
}
#[pymethod]
fn find(
&self,
@@ -905,15 +926,8 @@ impl PyString {
start: OptionalArg<Option<isize>>,
end: OptionalArg<Option<isize>>,
) -> isize {
let value = &self.value;
if let Some((start, end)) = adjust_indices(start, end, value.len()) {
match value[start..end].find(&sub.value) {
Some(num) => (start + num) as isize,
None => -1 as isize,
}
} else {
-1 as isize
}
self._find(sub, start, end, |r, s| r.find(s))
.map_or(-1, |v| v as isize)
}
#[pymethod]
@@ -923,15 +937,8 @@ impl PyString {
start: OptionalArg<Option<isize>>,
end: OptionalArg<Option<isize>>,
) -> isize {
let value = &self.value;
if let Some((start, end)) = adjust_indices(start, end, value.len()) {
match value[start..end].rfind(&sub.value) {
Some(num) => (start + num) as isize,
None => -1 as isize,
}
} else {
-1 as isize
}
self._find(sub, start, end, |r, s| r.rfind(s))
.map_or(-1, |v| v as isize)
}
#[pymethod]
@@ -942,15 +949,8 @@ impl PyString {
end: OptionalArg<Option<isize>>,
vm: &VirtualMachine,
) -> PyResult<usize> {
let value = &self.value;
if let Some((start, end)) = adjust_indices(start, end, value.len()) {
match value[start..end].find(&sub.value) {
Some(num) => Ok(start + num),
None => Err(vm.new_value_error("substring not found".to_owned())),
}
} else {
Err(vm.new_value_error("substring not found".to_owned()))
}
self._find(sub, start, end, |r, s| r.find(s))
.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
}
#[pymethod]
@@ -961,15 +961,8 @@ impl PyString {
end: OptionalArg<Option<isize>>,
vm: &VirtualMachine,
) -> PyResult<usize> {
let value = &self.value;
if let Some((start, end)) = adjust_indices(start, end, value.len()) {
match value[start..end].rfind(&sub.value) {
Some(num) => Ok(start + num),
None => Err(vm.new_value_error("substring not found".to_owned())),
}
} else {
Err(vm.new_value_error("substring not found".to_owned()))
}
self._find(sub, start, end, |r, s| r.rfind(s))
.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
}
#[pymethod]
@@ -1048,9 +1041,9 @@ impl PyString {
start: OptionalArg<Option<isize>>,
end: OptionalArg<Option<isize>>,
) -> usize {
let value = &self.value;
if let Some((start, end)) = adjust_indices(start, end, value.len()) {
self.value[start..end].matches(&sub.value).count()
let range = adjust_indices(start, end, self.value.len());
if range.is_normal() {
self.value[range].matches(&sub.value).count()
} else {
0
}
@@ -1764,12 +1757,22 @@ impl PySliceableSequence for String {
}
}
pub trait StringRange {
fn is_normal(&self) -> bool;
}
impl StringRange for std::ops::Range<usize> {
fn is_normal(&self) -> bool {
self.start <= self.end
}
}
// help get optional string indices
fn adjust_indices(
pub fn adjust_indices(
start: OptionalArg<Option<isize>>,
end: OptionalArg<Option<isize>>,
len: usize,
) -> Option<(usize, 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);
if end > len as isize {
@@ -1786,11 +1789,7 @@ fn adjust_indices(
start = 0;
}
}
if start > end {
None
} else {
Some((start as usize, end as usize))
}
start as usize..end as usize
}
// According to python following categories aren't printable: